Description

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Your KthLargest class will have a constructor which accepts an integer k and an integer array nums, which contains initial elements from the stream. For each call to the method KthLargest.add, return the element representing the kth largest element in the stream.

Example

int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3);   // returns 4
kthLargest.add(5);   // returns 5
kthLargest.add(10);  // returns 5
kthLargest.add(9);   // returns 8
kthLargest.add(4);   // returns 8

Note

You may assume that nums’ length ≥ k-1 and k ≥ 1.

Solution

此題要找出輸入流中第k大的數字。我們可以使用一個小頂堆，初始化時先用前k個數字將其填滿，由於堆頂元素是最小的，因此每新增一個數字時都將其和堆頂元素進行比較，只要比堆頂元素大就將堆頂元素移除，並把要新增的數字加入堆中。如此反覆，每次新增後堆中都儲存著當前輸入流中的前K大的數字，並且堆頂元素是其中最小的也就是第K大的數字。

class KthLargest {
	private Queue<Integer> heap;
	private int k;
	
    public KthLargest(int k, int[] nums) {
    	heap =
 
 new PriorityQueue<>(k);
    	this.k = k;
    	for(int num : nums)
    		add(num);
    }
    
    public int add(int val) {
    	if(heap.size() < k) {
    		heap.add(val);
    	} else if(val > heap.peek()) {
    		heap.poll();
    		heap.add(val);
    	}
    	return heap.peek();
    }
}