Top K Elements Pattern (Heap and QuickSelect)
1. Kth Largest Element (QuickSelect Algorithm)
| Approach | Data Structure | Time Complexity | Space | Best For |
|---|---|---|---|---|
| Min Heap (K Largest) | Size-K min heap | O(n log k) |
O(k) |
K much smaller than n |
| Max Heap (K Smallest) | Size-K max heap | O(n log k) |
O(k) |
K much smaller than n |
| QuickSelect | In-place partition | O(n) avg, O(n²) worst |
O(1) |
Single query, can modify array |
| Full Sort | Native sort | O(n log n) |
O(1) or O(n) |
Need all elements sorted anyway |
Example: Kth largest using min heap
class MinHeap {
constructor() {
this.heap = [];
}
push(val) {
this.heap.push(val);
this.bubbleUp(this.heap.length - 1);
}
pop() {
if (this.heap.length === 0) return null;
const min = this.heap[0];
const last = this.heap.pop();
if (this.heap.length > 0) {
this.heap[0] = last;
this.bubbleDown(0);
}
return min;
}
peek() {
return this.heap[0];
}
size() {
return this.heap.length;
}
bubbleUp(idx) {
while (idx > 0) {
const parent = Math.floor((idx - 1) / 2);
if (this.heap[idx] >= this.heap[parent]) break;
[this.heap[idx], this.heap[parent]] =
[this.heap[parent], this.heap[idx]];
idx = parent;
}
}
bubbleDown(idx) {
while (true) {
let smallest = idx;
const left = 2 * idx + 1;
const right = 2 * idx + 2;
if (left < this.heap.length &&
this.heap[left] < this.heap[smallest]) {
smallest = left;
}
if (right < this.heap.length &&
this.heap[right] < this.heap[smallest]) {
smallest = right;
}
if (smallest === idx) break;
[this.heap[idx], this.heap[smallest]] =
[this.heap[smallest], this.heap[idx]];
idx = smallest;
}
}
}
function findKthLargest(nums, k) {
const minHeap = new MinHeap();
for (const num of nums) {
minHeap.push(num);
if (minHeap.size() > k) {
minHeap.pop(); // Remove smallest
}
}
return minHeap.peek();
}
// [3,2,1,5,6,4], k=2 -> 5
Example: Kth largest using QuickSelect
function findKthLargest(nums, k) {
// Convert to kth smallest from end
return quickSelect(nums, 0, nums.length - 1,
nums.length - k);
}
function quickSelect(nums, left, right, k) {
if (left === right) return nums[left];
// Random pivot for better average case
const pivotIndex = left +
Math.floor(Math.random() * (right - left + 1));
const partitionIdx = partition(nums, left, right,
pivotIndex);
if (k === partitionIdx) {
return nums[k];
} else if (k < partitionIdx) {
return quickSelect(nums, left,
partitionIdx - 1, k);
} else {
return quickSelect(nums, partitionIdx + 1,
right, k);
}
}
function partition(nums, left, right, pivotIndex) {
const pivotValue = nums[pivotIndex];
// Move pivot to end
[nums[pivotIndex], nums[right]] =
[nums[right], nums[pivotIndex]];
let storeIndex = left;
for (let i = left; i < right; i++) {
if (nums[i] < pivotValue) {
[nums[storeIndex], nums[i]] =
[nums[i], nums[storeIndex]];
storeIndex++;
}
}
// Move pivot to final position
[nums[right], nums[storeIndex]] =
[nums[storeIndex], nums[right]];
return storeIndex;
}
// O(n) average, O(n²) worst case
2. K Closest Points to Origin
| Strategy | Implementation | Key Point | Optimization |
|---|---|---|---|
| Max Heap by Distance | Store K closest, evict farthest | Use max heap to track K smallest distances | Avoid sqrt - compare squared distances |
| QuickSelect on Distance | Partition by distance value | Find Kth closest, return all <= K | In-place, O(n) average |
| Custom Comparator | Define distance comparison function | Works with any distance metric | Reusable pattern |
| Distance Calculation | x² + y² (no sqrt needed) |
Squared distance preserves ordering | Faster computation |
Example: K closest points using max heap
class MaxHeap {
constructor(compareFn) {
this.heap = [];
this.compare = compareFn || ((a, b) => a - b);
}
push(val) {
this.heap.push(val);
this.bubbleUp(this.heap.length - 1);
}
pop() {
if (this.heap.length === 0) return null;
const max = this.heap[0];
const last = this.heap.pop();
if (this.heap.length > 0) {
this.heap[0] = last;
this.bubbleDown(0);
}
return max;
}
peek() {
return this.heap[0];
}
size() {
return this.heap.length;
}
bubbleUp(idx) {
while (idx > 0) {
const parent = Math.floor((idx - 1) / 2);
if (this.compare(this.heap[idx], this.heap[parent]) <= 0) break;
[this.heap[idx], this.heap[parent]] =
[this.heap[parent], this.heap[idx]];
idx = parent;
}
}
bubbleDown(idx) {
while (true) {
let largest = idx;
const left = 2 * idx + 1;
const right = 2 * idx + 2;
if (left < this.heap.length &&
this.compare(this.heap[left], this.heap[largest]) > 0) {
largest = left;
}
if (right < this.heap.length &&
this.compare(this.heap[right], this.heap[largest]) > 0) {
largest = right;
}
if (largest === idx) break;
[this.heap[idx], this.heap[largest]] =
[this.heap[largest], this.heap[idx]];
idx = largest;
}
}
}
function kClosest(points, k) {
const distanceSq = ([x, y]) => x * x + y * y;
const maxHeap = new MaxHeap((a, b) =>
distanceSq(a) - distanceSq(b));
for (const point of points) {
maxHeap.push(point);
if (maxHeap.size() > k) {
maxHeap.pop(); // Remove farthest
}
}
return maxHeap.heap;
}
// [[1,3],[-2,2]], k=1 -> [[-2,2]]
// Time: O(n log k), Space: O(k)
3. Top K Frequent Elements (HashMap + Heap)
| Approach | Steps | Complexity | When to Use |
|---|---|---|---|
| HashMap + Min Heap | Count frequencies, maintain K-size heap | O(n log k) |
K much smaller than unique elements |
| HashMap + Bucket Sort | Count, create frequency buckets | O(n) |
Optimal time, more space |
| HashMap + QuickSelect | Count, partition by frequency | O(n) avg |
One-time query, can modify |
| HashMap + Sort | Count, sort by frequency | O(n log n) |
Simple implementation |
Example: K frequent using min heap
function topKFrequent(nums, k) {
// Count frequencies
const freqMap = new Map();
for (const num of nums) {
freqMap.set(num, (freqMap.get(num) || 0) + 1);
}
// Min heap by frequency
const minHeap = new MinHeap();
for (const [num, freq] of freqMap.entries()) {
minHeap.push({ num, freq });
if (minHeap.size() > k) {
minHeap.pop();
}
}
return minHeap.heap.map(item => item.num);
}
// Custom MinHeap with frequency comparison
class MinHeap {
constructor() {
this.heap = [];
}
push(item) {
this.heap.push(item);
this.bubbleUp(this.heap.length - 1);
}
pop() {
if (this.heap.length === 0) return null;
const min = this.heap[0];
const last = this.heap.pop();
if (this.heap.length > 0) {
this.heap[0] = last;
this.bubbleDown(0);
}
return min;
}
size() { return this.heap.length; }
bubbleUp(idx) {
while (idx > 0) {
const parent = Math.floor((idx - 1) / 2);
if (this.heap[idx].freq >=
this.heap[parent].freq) break;
[this.heap[idx], this.heap[parent]] =
[this.heap[parent], this.heap[idx]];
idx = parent;
}
}
bubbleDown(idx) {
while (true) {
let smallest = idx;
const left = 2 * idx + 1;
const right = 2 * idx + 2;
if (left < this.heap.length &&
this.heap[left].freq <
this.heap[smallest].freq) {
smallest = left;
}
if (right < this.heap.length &&
this.heap[right].freq <
this.heap[smallest].freq) {
smallest = right;
}
if (smallest === idx) break;
[this.heap[idx], this.heap[smallest]] =
[this.heap[smallest], this.heap[idx]];
idx = smallest;
}
}
}
// [1,1,1,2,2,3], k=2 -> [1,2]
Example: K frequent using bucket sort
function topKFrequent(nums, k) {
// Count frequencies
const freqMap = new Map();
for (const num of nums) {
freqMap.set(num, (freqMap.get(num) || 0) + 1);
}
// Create buckets: index = frequency
const buckets = Array(nums.length + 1)
.fill(null).map(() => []);
for (const [num, freq] of freqMap.entries()) {
buckets[freq].push(num);
}
// Collect top K from highest frequencies
const result = [];
for (let i = buckets.length - 1; i >= 0; i--) {
for (const num of buckets[i]) {
result.push(num);
if (result.length === k) {
return result;
}
}
}
return result;
}
// Time: O(n), Space: O(n)
// Most efficient for large datasets
4. QuickSelect Partition Algorithm O(n) Average
| Concept | Details | Optimization | Trade-off |
|---|---|---|---|
| Partition Logic | Choose pivot, partition around it | Lomuto vs Hoare partition schemes | Lomuto simpler, Hoare fewer swaps |
| Pivot Selection | Random, median-of-3, or first element | Random gives O(n) expected time | Deterministic can be O(n²) worst case |
| Recursive Strategy | Only recurse on one half | T(n) = T(n/2) + O(n) = O(n) | Binary search-like pruning |
| In-place Modification | Modifies original array | O(1) extra space | Array order changed |
Example: QuickSelect with random pivot
function quickSelect(nums, k) {
// Find kth smallest (0-indexed)
return quickSelectHelper(nums, 0, nums.length - 1, k);
}
function quickSelectHelper(nums, left, right, k) {
if (left === right) return nums[left];
// Random pivot for average O(n)
const pivotIndex = left +
Math.floor(Math.random() * (right - left + 1));
const partitionIdx = partition(nums, left, right, pivotIndex);
if (k === partitionIdx) {
return nums[k];
} else if (k < partitionIdx) {
return quickSelectHelper(nums, left, partitionIdx - 1, k);
} else {
return quickSelectHelper(nums, partitionIdx + 1, right, k);
}
}
function partition(nums, left, right, pivotIndex) {
const pivotValue = nums[pivotIndex];
// Move pivot to end
[nums[pivotIndex], nums[right]] = [nums[right], nums[pivotIndex]];
let storeIndex = left;
// Lomuto partition scheme
for (let i = left; i < right; i++) {
if (nums[i] < pivotValue) {
[nums[storeIndex], nums[i]] = [nums[i], nums[storeIndex]];
storeIndex++;
}
}
// Move pivot to final position
[nums[right], nums[storeIndex]] = [nums[storeIndex], nums[right]];
return storeIndex;
}
// Average: O(n), Worst: O(n²)
// Space: O(1) if iterative, O(log n) recursion stack
Example: QuickSelect with median-of-3 pivot
function medianOf3Pivot(nums, left, right) {
const mid = Math.floor((left + right) / 2);
// Sort left, mid, right
if (nums[left] > nums[mid]) {
[nums[left], nums[mid]] = [nums[mid], nums[left]];
}
if (nums[mid] > nums[right]) {
[nums[mid], nums[right]] = [nums[right], nums[mid]];
}
if (nums[left] > nums[mid]) {
[nums[left], nums[mid]] = [nums[mid], nums[left]];
}
return mid; // Median is at mid
}
function quickSelectMedian3(nums, k) {
return helper(nums, 0, nums.length - 1, k);
function helper(nums, left, right, k) {
if (left === right) return nums[left];
const pivotIndex = medianOf3Pivot(nums, left, right);
const partitionIdx = partition(nums, left, right, pivotIndex);
if (k === partitionIdx) return nums[k];
else if (k < partitionIdx) {
return helper(nums, left, partitionIdx - 1, k);
} else {
return helper(nums, partitionIdx + 1, right, k);
}
}
}
// Better worst-case performance than random pivot
5. Merge K Sorted Lists (Min Heap)
| Scenario | Approach | Complexity | Application |
|---|---|---|---|
| K Sorted Lists | Min heap with (value, listIndex, elementIndex) | O(n log k) |
Merge K sorted arrays/streams |
| Smallest Range K Lists | Track max, min heap for min across lists | O(n log k) |
Find smallest range covering K lists |
| K Pairs Smallest Sum | Min heap with pair sums | O(k log k) |
Find K pairs with smallest sums |
| Stream Top K | Maintain size-K heap as stream flows | O(log k) per insert |
Real-time top K tracking |
Example: Merge K sorted lists
function mergeKLists(lists) {
const minHeap = new MinHeap((a, b) => a.value - b.value);
// Initialize heap with first element from each list
for (let i = 0; i < lists.length; i++) {
if (lists[i].length > 0) {
minHeap.push({
value: lists[i][0],
listIdx: i,
elemIdx: 0
});
}
}
const result = [];
while (minHeap.size() > 0) {
const { value, listIdx, elemIdx } = minHeap.pop();
result.push(value);
// Add next element from same list
if (elemIdx + 1 < lists[listIdx].length) {
minHeap.push({
value: lists[listIdx][elemIdx + 1],
listIdx,
elemIdx: elemIdx + 1
});
}
}
return result;
}
// [[1,4,5],[1,3,4],[2,6]]
// -> [1,1,2,3,4,4,5,6]
// Time: O(n log k), n = total elements, k = number of lists
Example: K pairs with smallest sums
function kSmallestPairs(nums1, nums2, k) {
if (!nums1.length || !nums2.length) return [];
const minHeap = new MinHeap((a, b) => a.sum - b.sum);
const result = [];
// Initialize with pairs (nums1[i], nums2[0])
for (let i = 0; i < Math.min(k, nums1.length); i++) {
minHeap.push({
sum: nums1[i] + nums2[0],
i,
j: 0
});
}
while (k > 0 && minHeap.size() > 0) {
const { sum, i, j } = minHeap.pop();
result.push([nums1[i], nums2[j]]);
k--;
// Add next pair from same nums1 index
if (j + 1 < nums2.length) {
minHeap.push({
sum: nums1[i] + nums2[j + 1],
i,
j: j + 1
});
}
}
return result;
}
// nums1=[1,7,11], nums2=[2,4,6], k=3
// -> [[1,2],[1,4],[1,6]]
6. Bucket Sort for Top K Elements
| Use Case | Implementation | Time | Constraint |
|---|---|---|---|
| Frequency-based Top K | Buckets indexed by frequency | O(n) |
Frequency range known: [0, n] |
| Score-based Top K | Buckets for score ranges | O(n + range) |
Discrete score values |
| Collect from Buckets | Iterate from highest to lowest | O(n) |
Stop when K elements collected |
| Advantage | Linear time, no heap overhead | Fastest for Top K | Only works with bounded integer keys |
Example: Top K frequent using bucket sort
function topKFrequentBucket(nums, k) {
// Step 1: Count frequencies
const freqMap = new Map();
for (const num of nums) {
freqMap.set(num, (freqMap.get(num) || 0) + 1);
}
// Step 2: Create buckets (index = frequency)
const buckets = Array(nums.length + 1).fill(null).map(() => []);
for (const [num, freq] of freqMap.entries()) {
buckets[freq].push(num);
}
// Step 3: Collect top K from highest frequencies
const result = [];
for (let freq = buckets.length - 1; freq >= 0 && result.length < k; freq--) {
for (const num of buckets[freq]) {
result.push(num);
if (result.length === k) {
return result;
}
}
}
return result;
}
// [1,1,1,2,2,3], k=2 -> [1,2]
// Time: O(n), Space: O(n)
// Key insight: Max frequency is n, so bucket array size is O(n)
Example: Sort characters by frequency using buckets
function frequencySort(s) {
// Count frequencies
const freqMap = new Map();
for (const char of s) {
freqMap.set(char, (freqMap.get(char) || 0) + 1);
}
// Create frequency buckets
const buckets = Array(s.length + 1).fill(null).map(() => []);
for (const [char, freq] of freqMap.entries()) {
buckets[freq].push(char);
}
// Build result from highest to lowest frequency
let result = '';
for (let freq = buckets.length - 1; freq >= 0; freq--) {
for (const char of buckets[freq]) {
result += char.repeat(freq);
}
}
return result;
}
// "tree" -> "eert" or "eetr"
// "cccaaa" -> "cccaaa" or "aaaccc"
Top K Elements Pattern Summary
- Heap Approach: O(n log k) time, O(k) space - best when K << n
- QuickSelect: O(n) average time, O(1) space - modifies array
- Bucket Sort: O(n) time, O(n) space - optimal for frequency problems
- Min Heap for Top K Largest: Keep K largest by evicting smallest
- Max Heap for Top K Smallest: Keep K smallest by evicting largest
- K Sorted Lists Merge: Use min heap with (value, listIdx, elemIdx)
Note: Choose data structure based on constraints: Use heap when K is much smaller than n, QuickSelect for one-time queries with modifiable arrays, and bucket sort when dealing with frequencies or bounded integer ranges (O(n) time optimal).
Warning: QuickSelect has O(n²) worst case with poor pivot selection. Use random pivot or median-of-3 for better average performance. Bucket sort only works when key range is bounded - won't work for arbitrary real numbers or unbounded ranges.