Kth Largest Element In a Stream

The problem

Report the kth largest value of a growing collection of numbers. You get operation names in operations and, at the same positions in values, each operation's arguments. "KthLargest" always comes first: its arguments are k and a list of starting numbers, and its result is null. Every other operation is "add": it adds its one number and gives the kth largest of the numbers so far. Equal numbers count separately: in 8, 5, 5 the second and third largest are both 5. Return one result per operation. The starting list may be empty or shorter than k, but after each add there are at least k numbers.

What to write

function solve(operations: string[], values: (number | number[])[][]): (number | null)[]

Input
operations = ["KthLargest","add","add","add","add","add"]
values = [[3,[4,5,8,2]],[3],[5],[10],[9],[4]]

Output
[null,4,5,5,8,8]

Input
operations = ["KthLargest","add","add"]
values = [[1,[]],[-3],[-2]]

Output
[null,-3,-2]

Constraints

Every add query has at least k values in total.

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Use the constraints

   Numbers are only added, never removed, and every add asks the same single question: the kth largest so far. You could sort everything again after each add, but most of that work repeats what the last sort already did.

2. Hint 2: Ask a smaller question

   Say k is 3 and you are holding 4, 5 and 8 when 2 arrives. Can 2 be the third largest now? Could it become the third largest after any later add? What about 4, if a 9 arrives next?

3. Hint 3: Name the pattern

   A question about the kth largest, asked again and again while values arrive, is the signal. You only need the k largest values and quick access to the smallest of them. A min-heap gives that: a container that always has its smallest value on top. Capped at k values, it is the heap of the best k. Read the lesson on heap of the best k.

4. Hint 4: State the invariant

   The heap always holds the k largest values seen so far, or every value while there are fewer than k. Equal values each take their own place. So after any add, the smallest value in the heap is the kth largest.

5. Hint 5: Sketch the pseudocode

   Read k and start an empty min-heap. To put a value in: push it, and if the heap now holds more than k values, pop the smallest. Put in every starting number this way. For each add: put its value in the same way, then answer with the smallest value in the heap. Answer null for KthLargest.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Approach

Keep a min-heap containing only the k largest values so far. Its root is the kth largest.

The slow way

Keep every number in a list. After each add, sort the list from largest to smallest and read the kth entry. It is correct, but every add sorts the whole collection again, and the list grows with every number that arrives.

The insight

The kth largest value is the smallest of the k largest values. A number outside those k can never come back, because numbers are only added. So keep only the k largest in a min-heap: a container that always has its smallest value on top, at its root. With k values in it, adding one or removing the smallest takes about log k steps.

Step by step

1. Keep a min-heap for the k largest values seen so far. The kth largest is the smallest of the k largest values. A min-heap keeps its smallest value on top, so the answer can be read in one step at any time.
2. Put the new value into the heap. The new value may or may not belong among the k largest. Putting it in first and trimming afterwards handles both cases the same way, with no comparison of your own.
3. If the heap now holds more than k values, remove the smallest. With k + 1 values in the heap, the smallest has k values that are at least as large. So the kth largest can be found without that value, now and after every later add.
4. Run every starting number through the same two steps. The starting numbers count exactly like added ones. Trimming while they go in leaves at most k values in the heap, however long the starting list is.
5. For each add, put the value in, then report the heap's smallest. After the value is in and the heap is trimmed, the heap holds exactly the k largest values so far, so its smallest is the kth largest. KthLargest reports nothing, so its result is null.

The same numbered steps are the comments in each solution.

function solve(operations, values) {
  // 1. Keep a min-heap for the k largest values seen so far
  const k = Number(values[0][0]);
  // JavaScript has no built-in heap, so this array is kept in heap order by
  // hand: no value is smaller than its parent, so the smallest is at index 0.
  // The parent of index i is at (i - 1) >>> 1, which is (i - 1) / 2 rounded
  // down, and its two children are at 2 * i + 1 and 2 * i + 2.
  const heap = [];
  function push(x) {
    // 2. Put the new value into the heap
    // It starts at the end and swaps with its parent while the parent is
    // larger.
    heap.push(x);
    let i = heap.length - 1;
    while (i) {
      const p = (i - 1) >>> 1;
      if (heap[p] <= heap[i]) break;
      [heap[p], heap[i]] = [heap[i], heap[p]];
      i = p;
    }
    // 3. If the heap now holds more than k values, remove the smallest
    if (heap.length > k) {
      // The last value takes the smallest value's place at index 0, then swaps
      // with its smaller child while that child is smaller.
      heap[0] = heap.pop();
      i = 0;
      while (true) {
        let child = 2 * i + 1;
        if (child >= heap.length) break;
        if (child + 1 < heap.length && heap[child + 1] < heap[child]) child++;
        if (heap[i] <= heap[child]) break;
        [heap[i], heap[child]] = [heap[child], heap[i]];
        i = child;
      }
    }
  }
  // 4. Run every starting number through the same two steps
  for (const x of values[0][1]) push(x);
  // 5. For each add, put the value in, then report the heap's smallest
  return operations.map((op, i) => {
    if (op === "KthLargest") return null;
    push(Number(values[i][0]));
    return heap[0];
  });
}

function solve(
  operations: string[],
  values: (number | number[])[][],
): (number | null)[] {
  // 1. Keep a min-heap for the k largest values seen so far
  const k = Number(values[0][0]);
  // TypeScript has no built-in heap, so this array is kept in heap order by
  // hand: no value is smaller than its parent, so the smallest is at index 0.
  // The parent of index i is at (i - 1) >>> 1, which is (i - 1) / 2 rounded
  // down, and its two children are at 2 * i + 1 and 2 * i + 2.
  const heap: number[] = [];
  function push(x: number): void {
    // 2. Put the new value into the heap
    // It starts at the end and swaps with its parent while the parent is
    // larger.
    heap.push(x);
    let i = heap.length - 1;
    while (i) {
      const p = (i - 1) >>> 1;
      if (heap[p] <= heap[i]) break;
      [heap[p], heap[i]] = [heap[i], heap[p]];
      i = p;
    }
    // 3. If the heap now holds more than k values, remove the smallest
    if (heap.length > k) {
      // The last value takes the smallest value's place at index 0, then swaps
      // with its smaller child while that child is smaller.
      heap[0] = heap.pop()!;
      i = 0;
      while (true) {
        let child = 2 * i + 1;
        if (child >= heap.length) break;
        if (child + 1 < heap.length && heap[child + 1] < heap[child]) child++;
        if (heap[i] <= heap[child]) break;
        [heap[i], heap[child]] = [heap[child], heap[i]];
        i = child;
      }
    }
  }
  // 4. Run every starting number through the same two steps
  for (const x of values[0][1] as number[]) push(x);
  // 5. For each add, put the value in, then report the heap's smallest
  return operations.map((op, i) => {
    if (op === "KthLargest") return null;
    push(Number(values[i][0]));
    return heap[0];
  });
}

import heapq


def solve(operations, values):
    # 1. Keep a min-heap for the k largest values seen so far
    k, nums = values[0]
    # heapq keeps a plain list in heap order: heap[0] is always its smallest value.
    heap = []

    def push(num):
        # 2. Put the new value into the heap
        heapq.heappush(heap, num)
        # 3. If the heap now holds more than k values, remove the smallest
        if len(heap) > k:
            heapq.heappop(heap)

    # 4. Run every starting number through the same two steps
    for num in nums:
        push(num)
    # 5. For each add, put the value in, then report the heap's smallest
    out = []
    for op, args in zip(operations, values):
        result = None
        if op == "add":
            push(args[0])
            result = heap[0]
        out.append(result)
    return out

Example 1, traced

Input
operations = ["KthLargest","add","add","add","add","add"]
values = [[3,[4,5,8,2]],[3],[5],[10],[9],[4]]

Output
[null,4,5,5,8,8]

Why it works

The heap always holds the k largest values seen so far, or all of them while there are fewer than k. Putting a value in and then removing the smallest of k + 1 keeps that true, because the removed value has k values at least as large. With exactly k values kept, the smallest one is the kth largest.

Time and space

Time O(log k) per add, extra space O(k). A heap has levels, each twice as wide as the one above, so k values fill about log k levels. Adding a value or removing the smallest moves one value up or down those levels: O(log k) per add. At most k + 1 values are stored: O(k) space.

n and m are input sizes, h is a tree's height, V and E are a graph's vertices and edges. The bounds describe this reference solution.

Watch out

• Put the new value in before you remove the smallest. If you make room first, a value smaller than everything kept still gets in: example 1 then answers 3 for the first add instead of 4.
• Trim while the starting numbers go in, not only inside add. Example 1 starts with four numbers but k is 3: left untrimmed, the heap stays one value too big and the first add answers 3 instead of 4.
• Report the smallest value in the heap, not the largest. The largest value kept is the largest number of all: example 1 would answer 8 for the first add instead of 4.

Pattern

Heap of the best k.

Keep going