请设计一个类型KthLargest,它每次从一个数据流中读取一个数字,并得出数据流已经读取的数字中第k(k≥1)大的数字。该类型的构造函数有两个参数:一个是整数k,另一个是包含数据流中最开始数字的整数数组nums(假设数组nums的长度大于k)。该类型还有一个函数add,用来添加数据流中的新数字并返回数据流中已经读取的数字的第k大数字。
例如,当k=3且nums为数组[4,5,8,2]时,调用构造函数创建类型KthLargest的实例之后,第1次调用add函数添加数字3,此时已经从数据流中读取了数字4、5、8、2和3,第3大的数字是4;第2次调用add函数添加数字5时,则返回第3大的数字5。
由于每次都需要找出k个数字中的最小值,因此可以把这k个数字保存到最小堆中。每当从数据流中读出一个数字,就先判断这个新的数字是不是有必要添加到最小堆中。如果最小堆中元素的数目还小于k,那么直接将它添加到最小堆中。如果最小堆中已经有k个元素,那么将其和位于堆顶的最小值进行比较。如果新读出的数字小于或等于堆中的最小值,那么堆中的k个数字都比它大,因此它不可能是k个最大的数字中的一个。由于只需要保存最大的k个数字,因此新读出的数字可以忽略。如果新的数字大于堆顶的数字,那么堆顶的数字就是第k+1大的数字,可以将它从堆中删除,并将新的数字添加到堆中,这样堆中保存的仍然是到目前为止从数据流中读出的最大的k个数字,此时第k大的数字正好位于最小堆的堆顶。
public class KthLargest {
private PriorityQueue<Integer> minHead;
private int size;
public KthLargest(int k, int[] nums) {
size = k;
minHead = new PriorityQueue<>();
for (int num : nums) {
add(num);
}
}
public static void main(String[] args) {
int[] nums = {4,5,8,2};
KthLargest kthLargest = new KthLargest(3, nums);
System.out.println(kthLargest.add(3));
System.out.println(kthLargest.add(5));
}
public int add(int val) {
if (minHead.size() < size) {
minHead.offer(val);
}
else if (val > minHead.peek()) {
minHead.poll();
minHead.offer(val);
}
return minHead.peek();
}
}