io.github.shanqiang.util.WindowUtil Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jstream Show documentation
The newest version!
package io.github.shanqiang.util;

import io.github.shanqiang.table.Row;

import java.util.ArrayList;
import java.util.List;
import java.util.TreeMap;

import static java.lang.Math.min;

public class WindowUtil {
    /**
     *
     * @param rows          all rows
     * @param columnName    compare by this column
     * @param n             N
     * @return              the top N indexes in rows compare by columnName column
     */
    public static int[] topN(List rows, String columnName, int n) {
        n = min(rows.size(), n);

        int size = 0;
        TreeMap> treeMap = new TreeMap<>();
        for (int i = 0; i < rows.size(); i++) {
            Row row = rows.get(i);
            Comparable comparable = row.getComparable(columnName);
            if (size < n) {
                putTreeMap(treeMap, comparable, i);
                size++;
            } else {
                Comparable leastKey = treeMap.firstKey();
                if (comparable.compareTo(leastKey) > 0) {
                    List values = treeMap.get(leastKey);
                    values.remove(values.size() - 1);
                    if (values.isEmpty()) {
                        treeMap.remove(leastKey);
                    }
                    putTreeMap(treeMap, comparable, i);
                }
            }
        }

        int[] ret = new int[n];
        int i = 1;
        for (List values : treeMap.values()) {
            for (Integer index : values) {
                ret[n - i] = index;
                i++;
            }
        }

        return ret;
    }

    private static void putTreeMap(TreeMap> treeMap, Comparable key, Integer value) {
        List values = treeMap.get(key);
        if (null == values) {
            values = new ArrayList<>();
            treeMap.put(key, values);
        }
        values.add(value);
    }
}