Problem 33 · Heaps

Kth Largest Element in an Array

MediumClassic DSA & systems

01Problem

Return the kth largest element of an unsorted array without fully sorting it.

02Key insight

Quickselect with a random pivot and three-way partition gives expected O(n); a size-k min-heap gives O(n log k) and works on streams. Know when to pick each.

03Solution

solution.pypython17 lines

import random

def find_kth_largest(nums: list[int], k: int) -> int:
    if not 1 <= k <= len(nums):
        raise ValueError("k out of range")
    candidates = nums
    while True:
        pivot = random.choice(candidates)
        greater = [x for x in candidates if x > pivot]
        equal_count = sum(1 for x in candidates if x == pivot)
        if k <= len(greater):
            candidates = greater
        elif k <= len(greater) + equal_count:
            return pivot
        else:
            k -= len(greater) + equal_count
            candidates = [x for x in candidates if x < pivot]

Try it first.

TimeO(n) expected, O(n^2) worst

SpaceO(n)

tests.py · 4 passing

tests.pypython4 lines

assert find_kth_largest([3, 2, 1, 5, 6, 4], 2) == 5
assert find_kth_largest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4) == 4
assert find_kth_largest([7, 7, 7, 7], 3) == 7
assert find_kth_largest([-1], 1) == -1

04Staff-level follow-up

Data is distributed across 1000 shards? Each shard returns its local top-k via a heap and a coordinator merges them (k*shards values); or binary-search on value with distributed count queries when k is huge.