Kth Largest Element in a Stream

Problem

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.

Implement KthLargest class:

• KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of integers nums.
• int add(int val) Appends the integer val to the stream and returns the element representing the kth largest element in the stream.

Class Skeleton:

class KthLargest {

    public KthLargest(int k, int[] nums) {
        
    }
    
    public int add(int val) {
        
    }
}

Examples

Example 1:

Input
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
Output
[null, 4, 5, 5, 8, 8]

Explanation
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

Solution

Method 1 - Using minHeap

What we can do is store the k largest elements in the heap. When we find element larger that root of minHeap, we will remove that element and add current element to heap.

Video Explanation

Here is the video explanation of the same: <div class="youtube-embed"><iframe src="https://www.youtube.com/embed/KA0KWgt3cns" frameborder="0" allowfullscreen></iframe></div>

Code

Java

class KthLargest {
	final PriorityQueue<Integer> minHeap;
	final int k;

	public KthLargest(int k, int[] nums) {
		this.k = k;
		minHeap = new PriorityQueue<>(k);

		for (int num: nums) {
			add(num);
		}
	}

	public int add(int val) {
		minHeap.offer(val);
		if (minHeap.size()>k) {
			minHeap.poll(val);
		}
		return minHeap.peek();
	}
}

Another way to write the add function:

	public int add(int val) {
		if (minHeap.size()<k) {
			minHeap.offer(val);
		} else if (minHeap.peek()<val) {
			minHeap.poll();
			minHeap.offer(val);
		}
		return minHeap.peek();
	}