TopK 的一种实现

【TopK】

直接上代码

/**
 * 如果需要更灵活的实现，比如传入 Compartor<? super E> 来进行初始化
 * 就需要在构造的时候，把比较器传入 PriorityQueue 的构造器
 */
public class TopK<E extends Comparable<E>> {

    private final int k;
    private final transient ReentrantLock lock = new ReentrantLock();
    /**
     * 默认为小顶堆
     */
    private final PriorityQueue<E> q = new PriorityQueue<>();

    /**
     * 构造一个 TopK 的实例
     *
     * @param k TopK 中的 k > 0
     */
    public TopK(int k) {
        if (k <= 0) throw new IllegalArgumentException("TopK k = " + k + " <= 0 !!!");
        this.k = k;
    }

    /**
     * 无脑往里面怼元素
     *
     * @param e 要放进去的元素
     */
    public void put(E e) {
        final ReentrantLock lock = this.lock;
        lock.lock();
        try {
            q.offer(e);
            if (q.size() > k) {
                // 加进去一个，总 size > k 了，就 poll 出去一个
                // poll 出去的由 q 保证一定是最小的一个，留下的就是 k 个最大的
                q.poll();
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 获取 TopK 的排序数据
     *
     * @return TopK 的排序数据
     */
    public List<E> get() {
        final ReentrantLock lock = this.lock;
        lock.lock();
        try {
            // 因为要重复使用，所以这里不可以往外 pop()
            ArrayList<E> es = new ArrayList<>(q);
            // 从大到小的排序
            es.sort(Comparator.reverseOrder());
            return es;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 清空所有数据
     */
    public void clear() {
        final ReentrantLock lock = this.lock;
        lock.lock();
        try {
            q.clear();
        } finally {
            lock.unlock();
        }
    }
}