Kth Largest Element in an Array — LeetCode 215 Heap vs Quickselect
Advertisement
Problem Statement
Given an integer array nums and an integer k, return the kth largest element. Note that you must find the kth largest in sorted order, not the kth distinct element.
Constraints:
- 1 <= k <= nums.length <= 10^5
- -10^4 <= nums[i] <= 10^4
Input: nums = [3, 2, 1, 5, 6, 4], k = 2
Output: 5Input: nums = [3, 2, 3, 1, 2, 4, 5, 5, 6], k = 4
Output: 4Why This Problem Matters
LeetCode 215 Kth Largest Element is one of the most-asked interview problems at Meta, Amazon, Google, and Bloomberg. It is a perfect "compare two paradigms" question: a heap gives an O(n log k) clean solution, while Quickselect averages O(n).
The choice signals seniority. Heap: simpler, deterministic, generalizes to streaming. Quickselect: optimal expected time, in-place, but worst-case O(n^2) unless you pick a randomized or median-of-medians pivot. Top candidates discuss both.
Keywords: "Kth largest interview", "Quickselect FAANG", "heap vs partition", "top K classic".
The Core Insight
The Kth largest is the smallest element in the top K. Either:
- Heap: maintain a min-heap of size K — root is the answer.
- Quickselect: partition like quicksort but recurse only into the side containing index
n - k.
For interviews, code the heap first (clean and obviously correct), then mention Quickselect as an optimization.
Visual Dry Run (Heap)
nums = [3, 2, 1, 5, 6, 4], k = 2.
| Step | Value | Heap (min, size <= 2) | After Pop |
|---|---|---|---|
| 1 | 3 | 3 | - |
| 2 | 2 | 2, 3 | - |
| 3 | 1 | 1, 3 (popped 2... actually 1) | 2, 3 |
| 4 | 5 | 3, 5 | - |
| 5 | 6 | 5, 6 | - |
| 6 | 4 | 5, 6 | - |
Top of heap = 5. Answer.
Solution (Heap)
import heapq
class Solution:
def findKthLargest(self, nums, k):
h = []
for v in nums:
heapq.heappush(h, v)
if len(h) > k:
heapq.heappop(h)
return h[0]class MinHeap {
constructor() { this.h = []; }
push(v) { this.h.push(v); this._up(this.h.length - 1); }
pop() {
const top = this.h[0], last = this.h.pop();
if (this.h.length) { this.h[0] = last; this._down(0); }
return top;
}
peek() { return this.h[0]; }
size() { return this.h.length; }
_up(i) {
while (i > 0) {
const p = (i - 1) >> 1;
if (this.h[p] <= this.h[i]) break;
[this.h[p], this.h[i]] = [this.h[i], this.h[p]];
i = p;
}
}
_down(i) {
const n = this.h.length;
while (true) {
const l = 2 * i + 1, r = 2 * i + 2;
let s = i;
if (l < n && this.h[l] < this.h[s]) s = l;
if (r < n && this.h[r] < this.h[s]) s = r;
if (s === i) break;
[this.h[s], this.h[i]] = [this.h[i], this.h[s]];
i = s;
}
}
}
var findKthLargest = function(nums, k) {
const h = new MinHeap();
for (const v of nums) {
h.push(v);
if (h.size() > k) h.pop();
}
return h.peek();
};Time: O(n log k). Space: O(k).
Solution (Quickselect)
import random
class Solution:
def findKthLargestQS(self, nums, k):
target = len(nums) - k
lo, hi = 0, len(nums) - 1
while lo <= hi:
p = self.partition(nums, lo, hi)
if p == target:
return nums[p]
if p < target:
lo = p + 1
else:
hi = p - 1
def partition(self, a, lo, hi):
pivot_idx = random.randint(lo, hi)
a[pivot_idx], a[hi] = a[hi], a[pivot_idx]
pivot = a[hi]
store = lo
for i in range(lo, hi):
if a[i] < pivot:
a[store], a[i] = a[i], a[store]
store += 1
a[store], a[hi] = a[hi], a[store]
return storeTime: O(n) average, O(n^2) worst. Space: O(1) extra.
Common Mistakes
- Using a max-heap of size N — wastes space and time vs size-K min-heap.
- Confusing kth largest with kth distinct largest.
- Off-by-one when targeting
len(nums) - kin Quickselect. - Forgetting random pivot — adversarial inputs trigger O(n^2).
- Sorting the array first — O(n log n) is allowed but not optimal.
Interview Tips
- Start with the heap solution, derive the time complexity, then offer Quickselect.
- Mention Hoare's selection algorithm and median-of-medians for guaranteed O(n).
- Discuss tradeoffs: heap is online and predictable, Quickselect is offline and amortized fast.
- For streams, the heap approach is the only option.
Follow-up Questions
- Stream of values? Use the size-K min-heap pattern (LeetCode 703).
- Need top K, not just the Kth? Heap stores all K; sort at the end if needed.
- Distributed across machines? Partial top-K per shard, merge with K-way merge.
- What if K = 1? Just take the max in O(n).
Key Takeaways
- LeetCode 215 has two optimal solutions: O(n log k) heap and O(n) average Quickselect.
- A min-heap of size K beats sorting for small K.
- Quickselect needs random or median-of-medians pivot to avoid O(n^2) worst case.
- The Kth largest is at index
n - kin the sorted array. - For streaming or unknown N, the heap is the right choice.
- Quickselect is in-place; heap uses O(k) extra space.
- Always mention both approaches in a senior interview.
Advertisement