K-th Largest Element in a Stream

K-th Largest Element in a Stream
Photo by Markus Spiske / Unsplash

The challenge is to find the kth largest element in a stream.

Solution:

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

Solution:

class KthLargest {

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

Explanation:
The idea is to use a priority queue (min-heap) with size k to our advantage. This data structure ensures that the queue will always be in a sorted order with the k-th element of the array being the minimum value at the head of the queue.

The algorithm keeps popping values from the top of the queue as and when it sees a bigger element added to the heap.

Complexity Analysis:

  • Heaps can retrieve the smallest or largest element in the queue at O(1) time
  • Addition of elements to the heap takes place in log(n) time.