java解决TopK问题

TopK问题描述

在不知道有多少元素，或者说元素个数不固定、很多时，普通的排序方式就不能再使用了，效率低下。
本文将使用java的api中的优先级队列，实现最大的topK和最小的topK。
也就是说求出大量数据中的前k个，可以是从大到小的topN，也可以是从小到大的topN。

java代码

package org.feng.queue;

import com.sun.istack.internal.NotNull;

import java.util.*;
import java.util.function.Consumer;

/**
 * Created by Feng on 2019/12/30 9:57
 * CurrentProject's name is java8
 * 使用 java 中的{@link PriorityQueue}来完成 TopK 计算
 * @see Comparable 是应该要存储的内容，保证有序
 * @see PriorityQueue 核心数据结构
 * @author Feng
 */
public class TopK<E extends Comparable> implements Iterable<E> {

    private PriorityQueue<E> priorityQueue;
    /**top的大小*/
    private int k;
    private Order order;

    /**
     * 迭代：拿到底层优先级队列的迭代器
     * @return Iterator
     * @see PriorityQueue
     */
    @Override
    public Iterator<E> iterator() {
        return priorityQueue.iterator();
    }

    @SuppressWarnings("unchecked")
    @Override
    public void forEach(Consumer action) {
        Objects.requireNonNull(action);
        for (E e: this) {
            action.accept(e);
        }
    }


    /**
     * 构造器：指定堆（优先级队列）的大小
     * @param k 大小
     */
    public TopK(int k){
        this(Order.MAX, k);
    }

    /**
     * 按照指定的 {@link Order}类型来确定是最大值的 TopK 还是最小值的 TopK
     * @param order {@code Order.MIN 表示最小；Order.MAX表示最大}
     * @param k 大小
     */
    public TopK(Order order, int k){
        this.k = k;
        this.order = order;
        if(Order.MIN.equals(order)){
            this.priorityQueue = new PriorityQueue<>(k, Collections.reverseOrder());
        } else {
            this.priorityQueue = new PriorityQueue<>(k);
        }
    }

    /**
     * 批量增加元素
     * @param collection 集合
     */
    public void addAll(Collection<? extends E> collection){
        if(Order.MIN.equals(order)){
            for (E e : collection) {
                addMin(e);
            }
            return;
        }

        for (E e : collection) {
            addMax(e);
        }
    }

    /**
     * 动态添加元素：tops 中始终是最大的 k 个元素
     * <p>
     *     1. 如果元素个数小于优先级队列的大小，则直接添加<br>
     *     2. 当元素个数不小于优先级队列大小时，先与队列中最小值比较，
     *          如果大于队列中最小值，则替换掉这个最小值；否则就不做操作。
     * </p>
     * @param e 元素，假定其实现了 {@link Comparable} 接口
     */
    @SuppressWarnings("unchecked")
    public void addMax(E e) {
        if(insert(e)){
            return;
        }
        E head = priorityQueue.peek();
        // 队列中的元素小于要增加的元素
        if(Objects.requireNonNull(head).compareTo(e) <= 0){
            // 新的元素替换掉原来的最小值，成为Topk中的值
            priorityQueue.poll();
            priorityQueue.add(e);
        }
    }


    /**
     * 最小值的 top 队列
     * @param e
     */
    @SuppressWarnings("unchecked")
    public void addMin(E e){
        if(insert(e)){
            return;
        }

        E head = priorityQueue.peek();
        if(Objects.requireNonNull(head).compareTo(e) > 0){
            // 新的元素替换掉原来的最大值，成为Topk中的值
            priorityQueue.poll();
            priorityQueue.add(e);
        }
    }

    /**
     * 队列未满时，直接插入数据
     * @param e 数据
     * @return 插入成功返回 true
     */
    public boolean insert(E e){
        if(priorityQueue.size() < k){
            return priorityQueue.add(e);
        }
        return false;
    }

    public <T>T[] toArray(T[] arr){
        return priorityQueue.toArray(arr);
    }

    public <T>T[] toArraySorted(T[] arr){
        T[] array = priorityQueue.toArray(arr);
        Arrays.sort(array);
        return array;
    }

    public <T>T[] toArraySorted(T[] arr, @NotNull Comparator<T> comparator){
        Objects.requireNonNull(comparator);
        T[] array = priorityQueue.toArray(arr);
        Arrays.sort(array, comparator);
        return array;
    }

    /**
     * 获取第 k 个最大或最小的元素
     * @return E
     */
    public E getKth(){
        return priorityQueue.peek();
    }

    @Override
    public String toString() {
        String orderString = Order.MIN.equals(order) ? "min" : "max";
        return "TopK[k = " + k + ", order = " + orderString + "] ,Tops" + priorityQueue.toString();
    }

    /**
     * 排序规则：正序和逆序
     */
    public enum Order{
        /**逆序：从小到大排序*/
        MIN,
        /**正序：从大到小排序*/
        MAX
    }
}

测试

package org.feng.queue;

import java.util.Arrays;
import java.util.Comparator;
import java.util.List;

/**
 * Created by Feng on 2019/12/30 10:23
 * CurrentProject's name is java8
 * @author Feng
 */
public class TopTest {
    public static void main(String[] args) {
        List<Integer> list = Arrays.asList(2, 3, 5, 11, 6, 77, 23, 43, 32, 8, 99);
        System.out.println("min:");
        // 设置大小为5
        TopK<Integer> minTop = new TopK<>(TopK.Order.MIN,5);
        // 添加元素
        minTop.addAll(list);
        // TopK[k = 5 order = reverse] ,Tops[8, 6, 3, 2, 5]
        System.out.println(minTop);
        minTop.forEach(System.out::println);

        System.out.println("max:");
        TopK<Integer> maxTop = new TopK<>(TopK.Order.MAX,5);
        maxTop.addAll(list);
        System.out.println(maxTop);
        maxTop.forEach(System.out::println);

        System.out.println("toArray:");
        Integer[] array = maxTop.toArray(new Integer[0]);
        System.out.println(Arrays.toString(array));

        System.out.println("toArraySorted1:");
        Integer[] sorted1 = maxTop.toArraySorted(new Integer[0]);
        System.out.println(Arrays.toString(sorted1));

        System.out.println("toArraySorted2:");
        Integer[] sorted2 = maxTop.toArraySorted(new Integer[0], Comparator.comparingInt(n -> n));
        System.out.println(Arrays.toString(sorted2));
    }
}

测试结果

控制台打印了：

min:
TopK[k = 5, order = min] ,Tops[8, 6, 3, 2, 5]
8
6
3
2
5
max:
TopK[k = 5, order = max] ,Tops[23, 32, 77, 43, 99]
23
32
77
43
99
toArray:
[23, 32, 77, 43, 99]
toArraySorted1:
[23, 32, 43, 77, 99]
toArraySorted2:
[23, 32, 43, 77, 99]

Process finished with exit code 0