Kth Largest Element in an Array

Given an integer array nums and an integer k, return the k-th largest value in the array.

This means the k-th largest by position in sorted order, not the k-th largest distinct value: in [6, 6, 3], both the first and second largest are 6.

Try to do better than sorting the whole array.

Sorting answers the question, but it does far more work than needed: it puts every element in place when you only care about one position. The k-th largest is the element at index n - k of the ascending sort, so the real question is "which value lands at index n - k?"

Quicksort's partition step answers part of that for free. After partitioning around a pivot, everything smaller is on the left and everything larger is on the right, so you know exactly which side contains index n - k. Recurse into that side only and ignore the other. On average each round throws away a constant fraction of the array, so the total work is n + n/2 + n/4 + … = O(n).

Two details keep it fast in practice: pick the pivot at random, so sorted input doesn't trigger the O(n²) worst case, and partition into three groups (less, equal, greater), so arrays full of duplicates don't either.

import random


def find_kth_largest(nums: list[int], k: int) -> int:
    target = len(nums) - k  # index of the answer in ascending order
    lo, hi = 0, len(nums) - 1
    while True:
        pivot = nums[random.randint(lo, hi)]
        # Partition nums[lo..hi] into < pivot, == pivot, > pivot.
        lt, i, gt = lo, lo, hi
        while i <= gt:
            if nums[i] < pivot:
                nums[lt], nums[i] = nums[i], nums[lt]
                lt += 1
                i += 1
            elif nums[i] > pivot:
                nums[i], nums[gt] = nums[gt], nums[i]
                gt -= 1
            else:
                i += 1
        if target < lt:
            hi = lt - 1
        elif target > gt:
            lo = gt + 1
        else:
            return pivot

function findKthLargest(nums, k) {
  const target = nums.length - k; // index of the answer in ascending order
  let lo = 0;
  let hi = nums.length - 1;
  while (true) {
    const pivot = nums[lo + Math.floor(Math.random() * (hi - lo + 1))];
    // Partition nums[lo..hi] into < pivot, == pivot, > pivot.
    let lt = lo;
    let i = lo;
    let gt = hi;
    while (i <= gt) {
      if (nums[i] < pivot) {
        [nums[lt], nums[i]] = [nums[i], nums[lt]];
        lt++;
        i++;
      } else if (nums[i] > pivot) {
        [nums[i], nums[gt]] = [nums[gt], nums[i]];
        gt--;
      } else {
        i++;
      }
    }
    if (target < lt) hi = lt - 1;
    else if (target > gt) lo = gt + 1;
    else return pivot;
  }
}

Give the sort answer first and say its cost, then offer the heap and quickselect as the two improvements, with their trade-offs: the heap is O(n log k), simple, and works on streams; quickselect is O(n) on average but mutates the array and has a quadratic worst case.

Interviewers usually want to hear why quickselect is linear on average (each step keeps roughly half), and how you defend against the worst case (random pivot, three-way partition). If you code quickselect, state the invariant of the partition loop before writing it; that's where the bugs are.