GroupBy.java

/* ---------------------------------------------------------------------
 * Numenta Platform for Intelligent Computing (NuPIC)
 * Copyright (C) 2014-2016, Numenta, Inc.  Unless you have an agreement
 * with Numenta, Inc., for a separate license for this software code, the
 * following terms and conditions apply:
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU Affero Public License for more details.
 *
 * You should have received a copy of the GNU Affero Public License
 * along with this program.  If not, see http://www.gnu.org/licenses.
 *
 * http://numenta.org/licenses/
 * ---------------------------------------------------------------------
 */
package org.numenta.nupic.util;

import java.util.List;
import java.util.function.Function;

import chaschev.lang.Pair;

/**
 * <p>
 * Allegory to the Python itertools.groupby. Objects of this class
 * take a list of inputs and a function to produce keys. The iterator
 * or foreach loop generates grouped return values based on the key
 * generated by the supplied function.<br>
 * For instance:<br>
 * <pre>
 * Given the list:
 * List<Integer> l = Arrays.asList(new Integer[] { 2, 4, 4, 5 });
 * and the function:
 * Function<Integer, Integer> lambda = x -> x * 3;
 * 
 * A GroupBy can be compose as such:
 * GroupBy<Integer, Integer> grouper = GroupBy.of(l, lambda);
 * 
 * ...then iterated over as such:
 * for(Pair<Integer, Integer> p : grouper) {
 *     System.out.println("Pair key: " + p.getKey() + ", pair value: " + p.getValue());
 * }
 * 
 * Outputs:
 * 
 * Pair key: 2, pair value: 6
 * Pair key: 4, pair value: 12
 * Pair key: 4, pair value: 12
 * Pair key: 5, pair value: 15
 * </p>
 * <p>
 * Notes: Read up on groupby here:
 *        https://docs.python.org/dev/library/itertools.html#itertools.groupby
 *        
 * </p>
 * @author cogmission
 *
 * @param <T>
 * @param <R>
 */
public class GroupBy<T, R> implements Generator<Pair<T, R>>, PeekableIterator<Pair<T, R>> {

    /** serial version */
    private static final long serialVersionUID = 1L;
    
    private List<T> iter;
    private Function<T, R> fn;
    private Generator<Integer> range;
    private Pair<T, R> next;
    
    
    /**
     * Constructs a new {@code GroupBy}
     * 
     * @param l     the {@link List} containing the items used as input to the
     *              key generating function.     
     * @param fn    the {@link Function} to be used to generate the keys which describe
     *              the like contents of each grouping.
     */
    public GroupBy(List<T> l, Function<T, R> fn) {
        this.iter = l;
        this.fn = fn;
        this.range = IntGenerator.of(0, iter.size());
        
        if(range.hasNext()) {
            T t = iter.get(range.next());
            next = new Pair<T, R>(t, fn.apply(t));
        }
    }
    
    /**
     * {@inheritDoc}
     */
    @Override
    public Pair<T, R> peek() {
        return next;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public boolean hasNext() {
        return next != null;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public Pair<T, R> next() {
        T t = range.hasNext() ? iter.get(range.next()) : null;
        Pair<T, R> ret = next;
        next = t != null ? new Pair<T, R>(t, fn.apply(t)) : null;
        return ret;
    }
    
    /**
     * Returns a new {@code GroupBy} composed from the specified list 
     * and key-generating {@link Function}
     * 
     * @param l     the {@link List} containing the items used as input to the
     *              key generating function.     
     * @param fn    the {@link Function} to be used to generate the keys which describe
     *              the like contents of each grouping.
     * @return
     */
    public static <T, R> GroupBy<T, R> of(List<T> l, Function<T, R> fn) {
        return new GroupBy<T, R>(l, fn);
    }
 }