Kth Largest Element in a Stream easy

Problem Statement

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.

Your KthLargest class will have a constructor that accepts an integer k and an integer array nums, which contains initial numbers from the stream. The class should also have a method add, which accepts a new number and updates the kth largest element in the stream.

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);   // returns 4
kthLargest.add(5);   // returns 5
kthLargest.add(10);  // returns 5
kthLargest.add(9);   // returns 8
kthLargest.add(4);   // returns 8

Example 2:

Input
["KthLargest", "add", "add"]
[[2, [0]], [-1], [1]]
Output
[null, 0, 1]

Explanation
KthLargest kthLargest = new KthLargest(2, [0]);
kthLargest.add(-1); // returns 0
kthLargest.add(1);  // returns 1

Steps:

1. Initialization: Use a min-heap (PriorityQueue in Java) of size k to store the k largest elements seen so far. During initialization, add the elements from the input nums array to the heap.

2. Adding a new element:

   • If the heap size is less than k, simply add the new element to the heap.
   • If the heap size is equal to k and the new element is greater than the smallest element in the heap (the root), remove the smallest element and add the new element.
   • Return the smallest element in the heap (which is the kth largest).

Explanation:

A min-heap is used because we want to efficiently track the smallest among the k largest elements. By keeping a heap of size k, the smallest element in the heap will always represent the kth largest element encountered so far. If a new element is larger than the current kth largest (smallest in the heap), it replaces it, maintaining the invariant.

Code (Java):

import java.util.PriorityQueue;

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

    public KthLargest(int k, int[] nums) {
        this.k = k;
        minHeap = new PriorityQueue<>(); // Min-heap by default

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

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

Complexity:

• Time complexity:

  • Constructor: O(N log k), where N is the length of nums. Adding N elements to a heap of size k takes at most O(N log k) time.
  • add method: O(log k) on average. Adding or removing an element from a heap of size k takes O(log k) time.

• Space complexity: O(k) to store the min-heap.