使用 PriorityQueue 求解 Top K 问题

package com.tomoyadeng.javabeginner.interview;

import java.util.*;
import java.util.stream.Stream;

public class TopK<T> {

    /**
     * 堆的边界，Top K 问题中的 K
     */
    private final int boundary;

    /**
     * 优先队列，用来构造一个有界的堆
     */
    private final PriorityQueue<T> boundaryHeap;


    /**
     * 通过自定义边界 boundary 可以求解 top K 问题
     * 通过自定义比较器 comparator 可以控制求解 top K 大 还是 top K 小
     * @param boundary 边界 K
     * @param comparator 数据比较器
     */
    public TopK(int boundary, Comparator<T> comparator) {
        this.boundary = boundary;
        boundaryHeap = new PriorityQueue<>(boundary, comparator);
    }

    /**
     * 求解数据流中的top K， 将结果写入List中
     * @param dataStream 数据流
     * @param results top K 结果
     */
    public void topK(Stream<T> dataStream, List<T> results) {
        dataStream.forEach(this::add);

        while (!boundaryHeap.isEmpty()) {
            results.add(boundaryHeap.poll());
        }
    }

    /**
     * 向有界堆中添加元素的帮助方法
     * @param t 待添加数据
     */
    private void add(T t) {
        boundaryHeap.add(t);
        if (boundaryHeap.size() > boundary) {
            boundaryHeap.poll();
        }
    }
}

public static void main(String[] args) {
    // 构造一个 范围为 [0, 2^30] 的 Integer 流，通过limit可以控制大小
    final int upLimit = 1 << 30;
    Stream<Integer> stream = Stream.generate(Math::random)
            .map(d -> d * upLimit)
            .map(d -> (int) Math.round(d))
            .limit(100_000_000);

    // 将 (o1, o2) -> (o1 - o2) 换成 (o1, o2) -> (o2 - o1) 可以求解 top K 小
    TopK<Integer> topK = new TopK<>(10, (o1, o2) -> (o1 - o2));

    List<Integer> results = new ArrayList<>();

    long startTime = System.currentTimeMillis();

    topK.topK(stream, results);

    long endTime = System.currentTimeMillis();

    System.out.println("results: " + results);
    System.out.println("cost: " + (endTime - startTime) / 1000.0);
}