~/problems / Heaps / Heaps and priority queues

Kth Largest Element in an Array

medium ~25 min

Write find_kth_largest(nums, k) that returns the value that would sit at position k (1-based) if nums were sorted from largest to smallest. Duplicates count separately: in [5, 5, 1] the 2nd largest is 5.

find_kth_largest([7, 2, 9, 4], 1)        # 9
find_kth_largest([7, 2, 9, 4], 3)        # 4
find_kth_largest([3, 3, 3, 1, 2], 4)     # 2

Constraints: 1 <= k <= len(nums) <= 200,000.

Aim for O(n log k) or average O(n), without sorting the whole list. Repeatedly finding and removing the maximum is O(n·k) and too slow when k is large. (Sorting passes the tests in Python, since the built-in sort is so fast, but it defeats the purpose of the drill.)

Show hint

You only ever need to remember the k largest values seen so far, and to know which of them is the smallest. Alternatively, partition around a pivot as in quicksort but only recurse into the side that holds the answer; watch out for many equal values.

Topic: Heaps and priority queues. heapq: top-k, k-way merge, two heaps for a running median.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc