Marauder Bot

Problem

Design a data structure initialized with an integer k and a starting list of numbers. It must support adding new numbers one at a time, and after each addition it should report the kth largest value seen so far among all numbers added, including the initial ones.

Example. With k = 2 and starting values [4, 5, 8, 2], adding 3 gives the numbers [4, 5, 8, 2, 3], whose second largest is 5, so the call returns 5.

Key idea

Re-sorting the entire collection after every addition works but does far more than necessary, since only the relative position of the kth largest value matters, not a full ordering. The structure only needs to remember the k largest values seen so far; anything smaller than the current kth largest is irrelevant to future answers as long as k stays fixed.

Maintain a min-heap capped at size k, holding the k largest numbers seen so far, with the smallest of those k sitting at the top. When a new number arrives, add it to the heap; if the heap now holds more than k elements, remove the smallest. The value left at the top of the heap after this adjustment is always the kth largest overall, because anything smaller was either never added or was evicted as no longer among the top k.

Solution

// Minimal binary heap; the comparator (hasPriority) decides min-heap or max-heap order.
class Heap<T> {
  private readonly items: T[] = [];

  constructor(private readonly hasPriority: (a: T, b: T) => boolean) {}

  get size(): number {
    return this.items.length;
  }

  peek(): T {
    return this.items[0];
  }

  push(item: T): void {
    this.items.push(item);
    let i = this.items.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.hasPriority(this.items[parent], this.items[i])) {
        break;
      }
      [this.items[parent], this.items[i]] = [this.items[i], this.items[parent]];
      i = parent;
    }
  }

  pop(): T {
    const top = this.items[0];
    const last = this.items.pop()!;
    if (this.items.length > 0) {
      this.items[0] = last;
      this.bubbleDown(0);
    }
    return top;
  }

  private bubbleDown(start: number): void {
    let i = start;
    while (true) {
      const left = 2 * i + 1;
      const right = 2 * i + 2;
      let next = i;
      if (left < this.items.length && this.hasPriority(this.items[left], this.items[next])) {
        next = left;
      }
      if (right < this.items.length && this.hasPriority(this.items[right], this.items[next])) {
        next = right;
      }
      if (next === i) {
        break;
      }
      [this.items[i], this.items[next]] = [this.items[next], this.items[i]];
      i = next;
    }
  }
}

class KthLargest {
  private readonly k: number;
  private readonly heap = new Heap<number>((a, b) => a < b); // min-heap: smallest of the top k sits on top

  constructor(k: number, nums: number[]) {
    this.k = k;
    for (const num of nums) {
      this.add(num);
    }
  }

  add(val: number): number {
    this.heap.push(val);
    if (this.heap.size > this.k) {
      this.heap.pop(); // evict the smallest once more than k values are held
    }
    return this.heap.peek(); // top is the kth largest seen so far
  }
}

Complexity

Time: O(log k) per addition, for the heap insert and possible removal; O(n log k) to seed the initial list of n numbers.
Space: O(k). The heap never holds more than k elements.

Watch out for

Use a min-heap, not a max-heap; the smallest of the top k is what needs to be evicted and inspected.
Do not rebuild the heap from scratch on each call; the whole benefit comes from updating it incrementally.

Pattern

This is the "top-k with a bounded heap" pattern: cap a heap at size k and let it self-prune to keep the k best elements, which generalizes directly to k closest points and similar running-statistics problems.

Kth Largest Element in a Stream