Kth Largest Element in a Stream Leetcode Solution

Problem Statement

In this problem, we have to design a class KthLargest() that initially has an integer k and an array of integers. We need to write a parameterized constructor for it when an integer k and array nums are passed as arguments. The class also has a function add(val) that adds a new element with value val into the stream of integers. For each add() call, we need to return an integer which is the Kth largest element in the running stream.

Example

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

4 5 5 8 8

Explanation:

Approach (Min Heaps)

Whenever it comes to finding the Kth smallest/largest element, max/min heaps almost every time serve the purpose. This is because of their flexible nature to hold elements in a sorted(prioritized) fashion. Here, we could have also used an array to contain the elements whenever a query is made. But, in order to find the Kth largest element in the array, we must have to use O(N) time for each query. Therefore, we can maintain a min-heap of size k, to find the kth largest element in O(1) time. Note that since we are using a min-heap, the topmost element would the smallest in the heap. And since we bound the heap size to be equal to k after every query, this top element would be the Kth largest in the overall stream(as the heap would keep k largest elements only).

Implementation of Kth Largest Element in a Stream Leetcode Solution

C++ Program

#include <bits/stdc++.h>

using namespace std;

class KthLargest {
public:
    priority_queue <int , vector <int> , greater <int> > pq;
    int K;

    KthLargest(int k, vector<int>& nums) {
        K = k;
        for(int &x : nums) {
            pq.push(x);
            if(pq.size() > k) {
                pq.pop();
            }
        }
    }

    int add(int val) {
        pq.push(val);
        if(pq.size() > K)
            pq.pop();
        return pq.top();
    }
};

int main() {
    vector <int> nums = {4 , 5 , 8 , 2};
    int k = 3;
    KthLargest stream(k , nums);
    cout << stream.add(3) << " ";
    cout << stream.add(5) << " ";
    cout << stream.add(10) << " ";
    cout << stream.add(9) << " ";
    cout << stream.add(4) << " ";
    cout << endl;
    return 0;
}

Java Program

import java.util.*;
import java.io.*;

class comp implements Comparator<Integer> {
    public int compare(Integer a , Integer b) {
        if(a > b)
            return 1;
        if(a < b)
            return -1;
        return 0;
    }
}

class KthLargest {
    int K;
    PriorityQueue <Integer> pq;

    public KthLargest(int k, int[] nums) {
        K = k;
        pq = new PriorityQueue <Integer> (new comp());
        for(int x : nums) {
            pq.add(x);
            if(pq.size() > k) {
                pq.remove();
            }
        }
    }

    int add(int val) {
        pq.add(val);
        if(pq.size() > K)
            pq.remove();
        return pq.peek();
    }
}

class KthLargestInStream {
    public static void main(String args[]) {
        int k = 3;
        int[] nums = {4 , 5 , 8 , 2};
        KthLargest stream = new KthLargest(k , nums);
        System.out.print(stream.add(3) + " ");
        System.out.print(stream.add(5) + " ");
        System.out.print(stream.add(10) + " ");
        System.out.print(stream.add(9) + " ");
        System.out.print(stream.add(4) + " ");
        System.out.println();
    }
}

4 5 5 8 8

Complexity Analysis of Kth Largest Element in a Stream Leetcode Solution

Time Complexity

O(N + Q), where N = size of the initial array (while calling the constructor). Q = Number of queries. This is because we traverse the array once and answer each query in O(1) time.

Space Complexity 

O(K), where K is the given argument (while calling the constructor). This is because we maintain the heap size to be exactly k, after any set of operations.