LeetCode#703: Kth Largest Element in a Stream

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();
    }
}