Heap: kth largest in a stream

Problem Statement:

You are given an infinite stream of number and a number k.

You need to return the kth largest element after every insert.

Example:

Input: arr = [1, 2, 9, 6, 3, 5, 7, 10] k = 3

Output: _ _ 1 2 3 5 6 7

Solution Explanation:

We will solve the problem by using min heap.

Insert first k elements into the heap.

Then after k inserts, if the current element is more than the rootm then pop the root and insert.

We will use heap because the heap root will always hold kth largest element

Time Complexity: O(1)
Space Complexity: O(1)

Code Solution

#include <iostream>
#include <vector>
#include <queue>
using namespace std;

void solution(vector<int> arr, int k)
{

   priority_queue<int, vector<int>, greater<int> > pq;

   for (int i=0; i < k-1; i++)
   {
      pq.push(arr[i]);
      cout << "_ ";
   }

   pq.push(arr[k-1]);

   for (int i = k; i < arr.size(); i++)
   {
        cout << pq.top() << " ";

       if (arr[i] > pq.top())
       {
           pq.pop();
           pq.push(arr[i]);
       }  
   }  

   cout << pq.top();
}

int main()
{
   vector<int> arr = {1, 2, 9, 6, 3, 5, 7, 10};
   int k = 3;

   solution(arr, k);
   
   return 0;
}

Output

_ _ 1 2 3 5 6 7