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.

Constraints

• 1 <= k <= 104
• 0 <= nums.length <= 104
• -104 <= nums[i] <= 104
• -104 <= val <= 104
• At most 104 calls will be made to add.
• It is guaranteed that there will be at least k elements in the array when you search for the kth element.

Solution

The problem Kth Largest Element in a Stream can be solved by implementing a priority queue with a limited capacity.

Implementation

static const int fast_io = []()
{
    std::ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    return 0;
}();

class KthLargest
{
  private:
    int capacity;
    priority_queue<int, vector<int>, greater<int>> pq;

  public:
    KthLargest(int k, vector<int> &nums)
    {
        capacity = k;
        for (int num : nums)
            add(num);
    }

    int add(int val)
    {
        if (pq.size() < capacity || pq.top() < val)
            pq.push(val);
        if (pq.size() > capacity)
            pq.pop();

        return pq.top();
    }
};

/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest* obj = new KthLargest(k, nums);
 * int param_1 = obj->add(val);
 */