Kth Largest Element In An Array

The problem

Return the kth largest number in nums: the number that would stand at position k, counting from 1, if the list were written from largest to smallest. Equal numbers each take their own place, so in 5, 5, 4 the second largest is 5 and the third largest is 4. k is never larger than the length of the list.

What to write

function solve(nums: number[], k: number): number

Input
nums = [3,2,1,5,6,4]
k = 2

Output
5

Input
nums = [3,2,3,1,2,4,5,5,6]
k = 4

Output
4

Constraints

1–10,000 integers; 1 ≤ k ≤ nums.length.

Hints

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

1. Hint 1: Use the constraints

   You are asked for one number, the kth largest, not for the whole list in order. Sorting all n numbers answers the question, but it also works out the order of every number below the kth, which nobody asked for.

2. Hint 2: Ask a smaller question

   Say k is 2 and you read 3, 2, 1, 5, 6, 4 from left to right. After the first three numbers, which of them could still be the answer in the end, and which is already out for good? What is the least you must remember as you go?

3. Hint 3: Name the pattern

   Being asked for the kth largest is the signal. Only the k largest numbers seen so far can still matter, and the one you need from them is the smallest. A min-heap, a container that always has its smallest value on top, capped at k values, does exactly this. This is the heap of the best k. Read the lesson on heap of the best k.

4. Hint 4: State the invariant

   After each number, the heap holds the k largest numbers read so far, or all of them while fewer than k have been read. When the last number has been read, the smallest value in the heap is the kth largest of the whole list.

5. Hint 5: Sketch the pseudocode

   Start an empty min-heap. For each number: push it, and if the heap now holds more than k values, pop the smallest. After the last number, return the smallest value in the heap, which is the one on top.

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 of the k largest numbers seen so far: add each number, and remove the smallest whenever the heap holds more than k. The smallest number left is the kth largest.

The slow way

Sort a copy of the list from largest to smallest and take the kth entry. It is correct and short, but it puts all n numbers in order, about n log n steps, when only the k largest matter, and it stores a second copy of the whole list.

The insight

The kth largest number is the smallest of the k largest numbers. So read the list once and keep only the k largest seen so far, in a min-heap: a container that always has its smallest value on top, at its root. Whenever it holds more than k values, the smallest one can go, because k others are at least as large.

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, where it can be read in one step once the whole list has gone through.
2. Put the next number into the heap. The number 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 it cannot be needed for the kth largest. Removing it keeps the heap at k values.
4. The smallest value left in the heap is the kth largest. After the last number the heap holds exactly the k largest numbers of the list, because the list has at least k numbers. The smallest of those k is the answer, and it is on top.

The same numbered steps are the comments in each solution.

function solve(nums, k) {
  // JavaScript has no built-in heap, so these two helpers keep an array in
  // min-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.
  function heapPush(heap, x) {
    // The new value 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;
    }
  }
  function heapPop(heap) {
    // The smallest value leaves. The last value takes its place at index 0,
    // then swaps with its smaller child while that child is smaller.
    const smallest = heap[0],
      last = heap.pop();
    if (!heap.length) return smallest;
    heap[0] = last;
    let 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;
    }
    return smallest;
  }
  // 1. Keep a min-heap for the k largest values seen so far
  const heap = [];
  for (const num of nums) {
    // 2. Put the next number into the heap
    heapPush(heap, num);
    // 3. If the heap now holds more than k values, remove the smallest
    if (heap.length > k) heapPop(heap);
  }
  // 4. The smallest value left in the heap is the kth largest
  return heap[0];
}

function solve(nums: number[], k: number): number {
  // TypeScript has no built-in heap, so these two helpers keep an array in
  // min-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.
  function heapPush(heap: number[], x: number): void {
    // The new value 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;
    }
  }
  function heapPop(heap: number[]): number {
    // The smallest value leaves. The last value takes its place at index 0,
    // then swaps with its smaller child while that child is smaller.
    const smallest = heap[0],
      last = heap.pop()!;
    if (!heap.length) return smallest;
    heap[0] = last;
    let 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;
    }
    return smallest;
  }
  // 1. Keep a min-heap for the k largest values seen so far
  const heap: number[] = [];
  for (const num of nums) {
    // 2. Put the next number into the heap
    heapPush(heap, num);
    // 3. If the heap now holds more than k values, remove the smallest
    if (heap.length > k) heapPop(heap);
  }
  // 4. The smallest value left in the heap is the kth largest
  return heap[0];
}

import heapq


def solve(nums, k):
    # 1. Keep a min-heap for the k largest values seen so far
    # heapq keeps a plain list in heap order: heap[0] is always its smallest value.
    heap = []
    for num in nums:
        # 2. Put the next number 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. The smallest value left in the heap is the kth largest
    return heap[0]

Example 2, traced

Input
nums = [3,2,3,1,2,4,5,5,6]
k = 4

Output
4

Why it works

The heap always holds the k largest numbers read so far, or all of them while fewer than k have been read. A number is removed only when k others in the heap are at least as large, so no number that is still needed is lost. At the end the heap holds the k largest of the list, and the smallest of them is the kth largest.

Time and space

Time O(n log k), extra space O(k). A heap has levels, each twice as wide as the one above, so k values fill about log k levels. Each of the n numbers costs one push and at most one removal along those levels: O(n log k). 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

• Do not remove repeated numbers first. Equal numbers each take their own place: with the repeats thrown away, example 2 returns 3 instead of 4.
• Remove the smallest value when the heap grows past k, not the largest. A heap with its largest value on top keeps the k smallest numbers and returns the largest of them: example 1 returns 2 instead of 5.
• Remove a value only when the heap holds more than k. Removing at k keeps k - 1 values: example 1 returns 6 instead of 5, and with k = 1 the heap ends up empty.

Pattern

Heap of the best k. Builds on Kth Largest Element In a Stream.

Keep going