Medium Quick Select Heaps
In plain English: Find the kth largest number in an unsorted array without fully sorting it.
Quick Select is partitioning (like quicksort) but only recursing into one half. Pick a pivot, partition around it, then check which partition the kth largest falls in. Average O(n), worst O(n^2).
Prompt
Given an integer array nums and an integer k, return the kth largest element. Solve in O(n) average time.
Try to write it from scratch before scrolling down.
Solution
import random
def find_kth_largest(nums, k):
target = len(nums) - k # convert to kth smallest (0-indexed)
def quickselect(lo, hi):
# random pivot avoids worst-case O(n^2) on sorted/adversarial inputs
pivot_idx = random.randint(lo, hi)
nums[pivot_idx], nums[hi] = nums[hi], nums[pivot_idx]
pivot = nums[hi]
store = lo # partition boundary: everything left of store is < pivot
for i in range(lo, hi):
if nums[i] < pivot:
nums[i], nums[store] = nums[store], nums[i]
store += 1
nums[store], nums[hi] = nums[hi], nums[store]
if store == target:
return nums[store]
elif store < target:
return quickselect(store + 1, hi)
else:
return quickselect(lo, store - 1)
return quickselect(0, len(nums) - 1)
# Test: find_kth_largest([3,2,1,5,6,4], 2) == 5
# Test: find_kth_largest([3,2,3,1,2,4,5,5,6], 4) == 4