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.