Kth Largest Element in a Stream

Difficulty: Easy

Design a class that keeps track of the kth largest element in a growing collection of numbers, as new numbers keep arriving one at a time.

The class is constructed with an integer k and an initial array of numbers. It supports one operation, add(val), which adds val to the collection and then returns the kth largest element in the collection so far (with duplicates counted - the kth largest, not the kth distinct value).

You're guaranteed there will always be at least k numbers in the collection whenever add is called.

Examples

Input: k = 3, nums = [4, 5, 8, 2]; then add(3), add(5), add(10), add(9), add(4)
Output: 4, 5, 5, 8, 8

Starting from [4,5,8,2], the 3 largest are 8,5,4, so the 3rd largest is 4. After add(3), the collection is [4,5,8,2,3] and the 3rd largest is still 4. After add(5) it becomes 5, and so on as more numbers arrive.

Constraints

  • 1 <= k <= 10^4

  • 0 <= nums.length <= 10^4

  • -10^4 <= nums[i] <= 10^4

  • At most 10^4 calls to add

Approach

Re-sorting everything on every call wastes time re-examining numbers that were never going to be in the top k anyway. Instead, keep a running collection of only the k largest values seen so far.

A min-heap of size k is perfect for this: it holds the k largest values, and its top (the minimum of those k) is exactly the kth largest value overall - because among the top k values, the smallest one is, by definition, the kth largest. When a new value arrives, add it to the heap; if that grows the heap past size k, remove the heap's minimum (it's no longer one of the top k). Whatever remains on top is the answer.

Solutions

Brute Force - Sort on Every Call

Keep all numbers in an array. Every time add() is called, push the new value, sort the whole array, and read off the kth largest by index.

class KthLargest {
  constructor(k, nums) {
    this.k = k;
    this.nums = [...nums];
  }

  add(val) {
    this.nums.push(val);
    this.nums.sort((a, b) => b - a);
    return this.nums[this.k - 1];
  }
}

Time: O(n log n) per add call, where n is the number of values seen so far · Space: O(n) to store every value ever added

Optimal - Min-Heap of Size k

Maintain a min-heap that only ever holds the k largest values seen so far. Its top is always the smallest of those k values - which is exactly the kth largest overall.

class MinHeap {
  constructor() {
    this.data = [];
  }

  size() {
    return this.data.length;
  }

  peek() {
    return this.data[0];
  }

  push(val) {
    this.data.push(val);
    let i = this.data.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.data[i] < this.data[parent]) {
        [this.data[i], this.data[parent]] = [this.data[parent], this.data[i]];
        i = parent;
      } else break;
    }
  }

  pop() {
    const top = this.data[0];
    const last = this.data.pop();
    if (this.data.length > 0) {
      this.data[0] = last;
      let i = 0;
      const n = this.data.length;
      while (true) {
        let smallest = i;
        const left = 2 * i + 1;
        const right = 2 * i + 2;
        if (left < n && this.data[left] < this.data[smallest]) smallest = left;
        if (right < n && this.data[right] < this.data[smallest]) smallest = right;
        if (smallest === i) break;
        [this.data[i], this.data[smallest]] = [this.data[smallest], this.data[i]];
        i = smallest;
      }
    }
    return top;
  }
}

class KthLargest {
  constructor(k, nums) {
    this.k = k;
    this.heap = new MinHeap();
    for (const num of nums) this.add(num);
  }

  add(val) {
    this.heap.push(val);
    if (this.heap.size() > this.k) this.heap.pop();
    return this.heap.peek();
  }
}

Time: O(log k) per add call · Space: O(k) - the heap never grows past size k