题目:
设计一个找到数据流中第 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
分析:
输入的数据是动态添加的,也就是说可以不断地从数据流中读取新的数据,数据流的数据量是无限的,该题如果数据存储在排序数组中,那么只需要O(1)的时间就能找出第k大的数字,但是采用这种方式缺点也是十分明显,如果从数据流中读取所有的数据都存到排序数组中,如果从数据流中读出n个数字,那么动态数组的大小为O(n),随着不断地从数据流中读出新的数据,空间复杂度可能会耗尽内存并且排序数组中添加新的数字时间复杂度不变,还是O(n)。
因此可以考虑数据结构堆,因为要找第k大的元素,因此使用小根堆,小根堆的堆顶元素就是这个堆的最小值,小根堆的元素个数最多为k,从头开始遍历数组,当小根堆的元素个数小于k的时候,入堆,当小根堆的元素个数大于等于k的时候并且当前遍历的值比小根堆的堆顶元素大则从堆中删除堆顶元素再添加当前元素,如此反复遍历结束后,堆中就剩下k个最大的值(遍历的数组中),而堆的堆顶元素则是整个堆最小的元素,所以总共就k个元素的堆中,并且这个堆里的元素是整个数组中最大的k个数,堆顶元素还是这k个数中最小的值,那么堆顶元素不就是整个数组第k大的数字嘛,因此返回小根堆的堆顶即可。
代码:
class KthLargest {
private PriorityQueue<Integer> minHeap;
private int size;
public KthLargest(int k, int[] nums) {
size = k;
minHeap = new PriorityQueue<>();
for (int num : nums) {
add(num);
}
}
public int add(int val) {
if (minHeap.size() < size){
minHeap.offer(val);
}else if (num > minHeap.peek()){
minHeap.poll();
minHeap.offer(val);
}
return minHeap.peek();
}
}