topK问题

发布时间 2023-03-27 00:50:49作者: zhengbiyu
设计一个找到数据流中第 k 大元素的类(class)。注意是排序后的第 k 大元素,不是第 k 个不同的元素。
请实现 KthLargest 类:
KthLargest(int k, int[] nums) 使用整数 k 和整数流 nums 初始化对象。
int add(int val) 将 val 插入数据流 nums 后,返回当前数据流中第 k 大的元素。 
示例:
输入:
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
输出:
[null, 4, 5, 5, 8, 8]
解释:
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3); // return 4
kthLargest.add(5); // return 5
kthLargest.add(10); // return 5
kthLargest.add(9); // return 8
kthLargest.add(4); // return 8
class KthLargest {
    PriorityQueue<Integer> pq;
    int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        pq = new PriorityQueue<Integer>();
        for (int x : nums) {
            add(x);
        }
    }
    
    public int add(int val) {
        pq.offer(val);
        if (pq.size() > k) {
            pq.poll();
        }
        return pq.peek();
    }
}
本题是我们求在一个数据流中的第 KK 大元素。所谓数据流,即是说我们写的算法需要支持 add() 函数;在力扣后台评测程序中会多次调用add()函数,每次调用都会向我们写的算法中添加一个元素。而题目要求的就是在每次 add() 之后,整个数据流(包括初始化的元素和所有 add 进来的元素)中的第 KK 大元素。
先说一个最暴力的解法:我们底层数据结构使用数组实现,当每次调用 add() 函数时,向数组中添加一个元素,然后调用 sort() 函数进行排序,返回排序后数组的第 KK 个数字。该做法在每次调用 add() 函数时的时间复杂度为 O(K*log(K))O(K∗log(K)) ,该时间复杂度太高,当 KK 很大 / add()调用次数太多的时候,一定会超时。
从上面的分析中,我们已经看出来了,使用数组的核心问题是:数组自身不带排序功能,只能用 sort() 函数,导致时间复杂度过高。
因此我们考虑使用自带排序功能的数据结构——堆。
在大根堆(图一)中,父节点的值比每一个子节点的值都要大。在小根堆(图二)中,父节点的值比每一个子节点的值都要小。
本题的操作步骤如下:
使用大小为 KK 的小根堆,在初始化的时候,保证堆中的元素个数不超过 KK 。
在每次 add() 的时候,将新元素 push() 到堆中,如果此时堆中的元素超过了 KK,那么需要把堆中的最小元素(堆顶)pop() 出来。
此时堆中的最小元素(堆顶)就是整个数据流中的第 KK 大元素。
问答:
为什么使用小根堆?
因为我们需要在堆中保留数据流中的前 KK 大元素,使用小根堆能保证每次调用堆的 pop() 函数时,从堆中删除的是堆中的最小的元素(堆顶)。
为什么能保证堆顶元素是第 KK 大元素?
因为小根堆中保留的一直是堆中的前 KK 大的元素,堆的大小是 KK,所以堆顶元素是第 KK 大元素。
每次 add() 的时间复杂度是多少?
每次 add() 时,调用了堆的 push() 和 pop() 方法,两个操作的时间复杂度都是 log(K)log(K).