703. Kth Largest Element In a Stream - Explanation

Description

Design a class to find the kth largest integer in a stream of values, including duplicates. E.g. the 2nd largest from [1, 2, 3, 3] is 3. The stream is not necessarily sorted.

Implement the following methods:

constructor(int k, int[] nums) Initializes the object given an integer k and the stream of integers nums.
int add(int val) Adds the integer val to the stream and returns the kth largest integer in the stream.

Example 1:

Input:
["KthLargest", [3, [1, 2, 3, 3]], "add", [3], "add", [5], "add", [6], "add", [7], "add", [8]]

Output:
[null, 3, 3, 3, 5, 6]

Explanation:
KthLargest kthLargest = new KthLargest(3, [1, 2, 3, 3]);
kthLargest.add(3);   // return 3
kthLargest.add(5);   // return 3
kthLargest.add(6);   // return 3
kthLargest.add(7);   // return 5
kthLargest.add(8);   // return 6

Constraints:

1 <= k <= 1000
0 <= nums.length <= 1000
-1000 <= nums[i] <= 1000
-1000 <= val <= 1000
There will always be at least k integers in the stream when you search for the kth integer.

Topics

Recommended Time & Space Complexity

You should aim for a solution with O(mlogk) time and O(k) space, where m is the number of times add() is called, and k represents the rank of the largest number to be tracked (i.e., the k-th largest element).

Hint 1

A brute force solution would involve sorting the array in every time a number is added using add(), and then returning the k-th largest element. This would take O(m * nlogn) time, where m is the number of calls to add() and n is the total number of elements added. However, do we really need to track all the elements added, given that we only need the k-th largest element? Maybe you should think of a data structure which can maintain only the top k largest elements.

Hint 2

We can use a Min-Heap, which stores elements and keeps the smallest element at its top. When we add an element to the Min-Heap, it takes O(logk) time since we are storing k elements in it. Retrieving the top element (the smallest in the heap) takes O(1) time. How can this be useful for finding the k-th largest element?

Hint 3

The k-th largest element is the smallest element among the top k largest elements. This means we only need to maintain k elements in our Min-Heap to efficiently determine the k-th largest element. Whenever the size of the Min-Heap exceeds k, we remove the smallest element by popping from the heap. How do you implement this?

Hint 4

We initialize a Min-Heap with the elements of the input array. When the add() function is called, we insert the new element into the heap. If the heap size exceeds k, we remove the smallest element (the root of the heap). Finally, the top element of the heap represents the k-th largest element and is returned.

Company Tags

Google2 Meta2 Amazon2 Microsoft5 Tinder4 Bloomberg3 Adobe2 Goldman Sachs2 Wells Fargo2

Prerequisites

Before attempting this problem, you should be comfortable with:

Heap Data Structure - Understanding min-heaps and how they maintain the smallest element at the root
Priority Queues - Using built-in heap implementations for efficient insertion and extraction
Sorting Algorithms - Understanding basic sorting as a baseline approach

1. Sorting

Intuition

We want the k-th largest number in a stream of values.
The simplest approach:
Every time a new value comes in, insert it, sort the list, and then pick the element at position len(arr) - k.

Sorting keeps the numbers in increasing order, so the k-th largest element will always sit at the same index.
This method is easy to understand but slow because sorting happens every time add() is called.

Algorithm

Initialization

Store k.
Store the initial numbers in an array.

add(val)

Append val to the array.
Sort the array.
Return the element at index len(arr) - k (the k-th largest).

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.arr = nums

    def add(self, val: int) -> int:
        self.arr.append(val)
        self.arr.sort()
        return self.arr[len(self.arr) - self.k]

class KthLargest {
    List<Integer> arr;
    int K;
    public KthLargest(int k, int[] nums) {
        K = k;
        arr = new ArrayList();
        for (int i = 0; i < nums.length; i++) {
            arr.add(nums[i]);
        }
    }

    public int add(int val) {
        arr.add(val);
        Collections.sort(arr);
        return arr.get(arr.size() - K);
    }
}

class KthLargest {
public:
    vector<int> arr;
    int k;
    KthLargest(int k, vector<int>& nums) {
        this->arr = nums;
        this->k = k;
    }

    int add(int val) {
        arr.push_back(val);
        sort(arr.begin(), arr.end());
        return arr[arr.size() - k];
    }
};

class KthLargest {
    /**
     * @param {number} k
     * @param {number[]} nums
     */
    constructor(k, nums) {
        this.arr = nums;
        this.k = k;
    }

    /**
     * @param {number} val
     * @return {number}
     */
    add(val) {
        this.arr.push(val);
        this.arr.sort((a, b) => a - b);
        return this.arr[this.arr.length - this.k];
    }
}

public class KthLargest {
    private List<int> arr;
    private int K;

    public KthLargest(int k, int[] nums) {
        arr = new List<int>(nums);
        K = k;
    }

    public int Add(int val) {
        arr.Add(val);
        arr.Sort();
        return arr[arr.Count - K];
    }
}

type KthLargest struct {
    k   int
    arr []int
}

func Constructor(k int, nums []int) KthLargest {
    return KthLargest{k: k, arr: nums}
}

func (this *KthLargest) Add(val int) int {
    this.arr = append(this.arr, val)
    sort.Ints(this.arr)
    return this.arr[len(this.arr)-this.k]
}

class KthLargest(k: Int, nums: IntArray) {
    private val k = k
    private val arr = nums.toMutableList()

    fun add(`val`: Int): Int {
        arr.add(`val`)
        arr.sort()
        return arr[arr.size - k]
    }
}

class KthLargest {
    private var k: Int
    private var arr: [Int]

    init(_ k: Int, _ nums: [Int]) {
        self.k = k
        self.arr = nums
    }

    func add(_ val: Int) -> Int {
        arr.append(val)
        arr.sort()
        return arr[arr.count - k]
    }
}

Time & Space Complexity

Time complexity: $O(m * n\log n)$
Space complexity:
- $O(m)$ extra space.
- $O(1)$ or $O(n)$ space depending on the sorting algorithm.

Where $m$ is the number of calls made to $add()$ and $n$ is the current size of the array.

2. Min-Heap

Intuition

To maintain the k-th largest element in a stream of numbers, we do not need to store all values.
Instead, we only need to keep track of the k largest elements seen so far.

A min-heap of size k is perfect for this:

A min-heap always keeps the smallest value at the top.
If the heap contains the k largest elements,
then the smallest among them is exactly the k-th largest overall.
Whenever a new number arrives:
- If we add it and the heap grows beyond k,
  we remove the smallest element - because it cannot be in the top k anymore.

This way, the heap always holds exactly the top k elements, and retrieving the k-th largest is O(1).

Algorithm

Initialization

Insert all initial numbers into a min-heap.
If the heap size becomes greater than k, repeatedly remove the smallest element.
- After this, the heap contains exactly k elements.

add(value)

Insert the new value into the min-heap.
If heap size > k:
- Remove the smallest element (the heap root).
Return the heap's smallest element (the root), which is now the k-th largest.

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.minHeap, self.k = nums, k
        heapq.heapify(self.minHeap)
        while len(self.minHeap) > k:
            heapq.heappop(self.minHeap)

    def add(self, val: int) -> int:
        heapq.heappush(self.minHeap, val)
        if len(self.minHeap) > self.k:
            heapq.heappop(self.minHeap)
        return self.minHeap[0]

class KthLargest {

    private PriorityQueue<Integer> minHeap;
    private int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        this.minHeap = new PriorityQueue<>();
        for (int num : nums) {
            minHeap.offer(num);
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }
    }

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

class KthLargest {
private:
    priority_queue<int, vector<int>, greater<int>> minHeap;
    int k;

public:
    KthLargest(int k, vector<int>& nums) {
        this->k = k;
        for (int num : nums) {
            minHeap.push(num);
            if (minHeap.size() > k) {
                minHeap.pop();
            }
        }
    }

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

/**
 * const { MinPriorityQueue } = require('@datastructures-js/priority-queue');
 */

class KthLargest {
    /**
     * @param {number} k
     * @param {number[]} nums
     */
    constructor(k, nums) {
        this.minHeap = new MinPriorityQueue();
        this.k = k;

        for (const num of nums) {
            this.minHeap.enqueue(num);
        }

        while (this.minHeap.size() > k) {
            this.minHeap.dequeue();
        }
    }

    /**
     * @param {number} val
     * @return {number}
     */
    add(val) {
        this.minHeap.enqueue(val);
        if (this.minHeap.size() > this.k) {
            this.minHeap.dequeue();
        }
        return this.minHeap.front();
    }
}

public class KthLargest {

    private PriorityQueue<int, int> minHeap;
    private int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        this.minHeap = new PriorityQueue<int, int>();
        foreach (int num in nums) {
            minHeap.Enqueue(num, num);
            if (minHeap.Count > k) {
                minHeap.Dequeue();
            }
        }
    }

    public int Add(int val) {
        minHeap.Enqueue(val, val);
        if (minHeap.Count > k) {
            minHeap.Dequeue();
        }
        return minHeap.Peek();
    }
}

type KthLargest struct {
    minHeap *priorityqueue.Queue
    k       int
}

func Constructor(k int, nums []int) KthLargest {
    minHeap := priorityqueue.NewWith(utils.IntComparator)
    for _, num := range nums {
        minHeap.Enqueue(num)
    }
    for minHeap.Size() > k {
        minHeap.Dequeue()
    }
    return KthLargest{minHeap: minHeap, k: k}
}

func (this *KthLargest) Add(val int) int {
    this.minHeap.Enqueue(val)
    if this.minHeap.Size() > this.k {
        this.minHeap.Dequeue()
    }
    top, _ := this.minHeap.Peek()
    return top.(int)
}

class KthLargest(k: Int, nums: IntArray) {
    private val k = k
    private val minHeap = PriorityQueue<Int>()

    init {
        for (num in nums) {
            minHeap.offer(num)
        }
        while (minHeap.size > k) {
            minHeap.poll()
        }
    }

    fun add(`val`: Int): Int {
        minHeap.offer(`val`)
        if (minHeap.size > k) {
            minHeap.poll()
        }
        return minHeap.peek()
    }
}

class KthLargest {
    private var minHeap: Heap<Int>
    private let k: Int

    init(_ k: Int, _ nums: [Int]) {
        self.k = k
        self.minHeap = Heap<Int>()
        for num in nums {
            minHeap.insert(num)
            if minHeap.count > k {
                minHeap.popMin()
            }
        }
    }

    func add(_ val: Int) -> Int {
        minHeap.insert(val)
        if minHeap.count > k {
            minHeap.popMin()
        }
        return minHeap.min!
    }
}

Time & Space Complexity

Time complexity: $O(m * \log k)$
Space complexity: $O(k)$

Where $m$ is the number of calls made to $add()$ .

Common Pitfalls

Using a Max-Heap Instead of Min-Heap

A common mistake is using a max-heap to find the k-th largest element. While it seems intuitive to keep the largest elements at the top, a max-heap would require storing all elements and repeatedly extracting the maximum k times for each query. A min-heap of size k is the correct choice because the root always holds the k-th largest element directly, allowing O(1) retrieval after each insertion.

Forgetting to Maintain Heap Size

When using a min-heap, it is essential to remove the smallest element whenever the heap size exceeds k. Failing to do so means the heap grows unbounded, and the root no longer represents the k-th largest element. Always check and pop after each insertion to keep exactly k elements in the heap.