Top K Elements Pattern — Heaps, Quickselect and Bucket Sort

Any question containing "the k largest", "the k most frequent", "the k closest" is this pattern. Sorting answers all of them in O(n log n) and will be accepted — the interview is about knowing that O(n log k) and sometimes O(n) are available, and why.

The counter-intuitive core

To keep the K largest elements, use a MIN-heap of size K.

That inversion trips people up every time, and the reasoning is short: the heap's root is the worst thing you are currently keeping. Each new element only needs comparing against that one value — better, so evict the root and push; worse, so discard. A max-heap would put the best element at the root, which tells you nothing about what to throw away.

import heapq

def k_largest(nums, k):
    heap = []
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)      # drop the smallest we are holding
    return heap

List<Integer> kLargest(int[] nums, int k) {
    // A MIN-heap for the K LARGEST: the root is the worst thing we are
    // keeping, so it is the only value a newcomer must beat.
    PriorityQueue<Integer> heap = new PriorityQueue<>();

    for (int x : nums) {
        heap.offer(x);
        if (heap.size() > k) heap.poll();   // drop the smallest we hold
    }
    return new ArrayList<>(heap);           // O(n log k), O(k) space
}

vector<int> kLargest(const vector<int>& nums, int k) {
    // A MIN-heap for the K LARGEST — greater<int> makes it a min-heap.
    priority_queue<int, vector<int>, greater<int>> heap;

    for (int x : nums) {
        heap.push(x);
        if ((int)heap.size() > k) heap.pop();   // drop the smallest we hold
    }

    vector<int> out;
    while (!heap.empty()) { out.push_back(heap.top()); heap.pop(); }
    return out;
}

function kLargest(nums, k) {
  // JavaScript has no built-in heap. For a one-shot query, a partial
  // selection sort over k is fine; for a stream, implement a binary heap.
  const heap = [];   // kept sorted ascending, so heap[0] is the smallest

  for (const x of nums) {
    if (heap.length < k) {
      heap.push(x);
      heap.sort((a, b) => a - b);
    } else if (x > heap[0]) {
      heap[0] = x;                    // evict the worst we are keeping
      heap.sort((a, b) => a - b);
    }
  }
  return heap;
}

The heap never exceeds k entries, so each push and pop is O(log k) and the whole scan is O(n log k) with O(k) space. On a stream, or on n in the millions with k in the tens, that gap is the answer.

Python's heapq is a min-heap only. For a max-heap, push negated values — and remember to negate back on the way out.

Frequency questions

Top K Frequent Elements adds a counting pass, then applies the same heap to the

(count, value) pairs:

from collections import Counter

def top_k_frequent(nums, k):
    counts = Counter(nums)
    return heapq.nlargest(k, counts.keys(), key=counts.get)

List<Integer> topKFrequent(int[] nums, int k) {
    Map<Integer, Integer> counts = new HashMap<>();
    for (int x : nums) counts.merge(x, 1, Integer::sum);

    // Min-heap on the count, capped at k.
    PriorityQueue<Integer> heap =
        new PriorityQueue<>(Comparator.comparingInt(counts::get));

    for (int value : counts.keySet()) {
        heap.offer(value);
        if (heap.size() > k) heap.poll();
    }
    return new ArrayList<>(heap);
}

vector<int> topKFrequent(const vector<int>& nums, int k) {
    unordered_map<int, int> counts;
    for (int x : nums) counts[x]++;

    // Min-heap on the count, capped at k.
    auto cmp = [](const pair<int,int>& a, const pair<int,int>& b) {
        return a.second > b.second;
    };
    priority_queue<pair<int,int>, vector<pair<int,int>>, decltype(cmp)> heap(cmp);

    for (const auto& entry : counts) {
        heap.push(entry);
        if ((int)heap.size() > k) heap.pop();
    }

    vector<int> out;
    while (!heap.empty()) { out.push_back(heap.top().first); heap.pop(); }
    return out;
}

function topKFrequent(nums, k) {
  const counts = new Map();
  for (const x of nums) counts.set(x, (counts.get(x) ?? 0) + 1);

  return [...counts.entries()]
    .sort((a, b) => b[1] - a[1])
    .slice(0, k)
    .map(([value]) => value);
}

But this problem has a better answer, and it is the one to reach for: bucket sort. A frequency can never exceed n, so index an array by count and walk it downward.

def top_k_frequent_buckets(nums, k):
    counts = Counter(nums)
    buckets = [[] for _ in range(len(nums) + 1)]
    for value, count in counts.items():
        buckets[count].append(value)

    out = []
    for count in range(len(buckets) - 1, 0, -1):
        for value in buckets[count]:
            out.append(value)
            if len(out) == k:
                return out
    return out

List<Integer> topKFrequentBuckets(int[] nums, int k) {
    Map<Integer, Integer> counts = new HashMap<>();
    for (int x : nums) counts.merge(x, 1, Integer::sum);

    // A frequency can never exceed n, so the bucket array is bounded.
    List<List<Integer>> buckets = new ArrayList<>();
    for (int i = 0; i <= nums.length; i++) buckets.add(new ArrayList<>());
    for (var e : counts.entrySet()) buckets.get(e.getValue()).add(e.getKey());

    List<Integer> out = new ArrayList<>();
    for (int count = nums.length; count > 0 && out.size() < k; count--) {
        for (int value : buckets.get(count)) {
            out.add(value);
            if (out.size() == k) return out;
        }
    }
    return out;   // O(n) — no log factor at all
}

vector<int> topKFrequentBuckets(const vector<int>& nums, int k) {
    unordered_map<int, int> counts;
    for (int x : nums) counts[x]++;

    // A frequency can never exceed n, so the bucket array is bounded.
    vector<vector<int>> buckets(nums.size() + 1);
    for (const auto& e : counts) buckets[e.second].push_back(e.first);

    vector<int> out;
    for (int count = (int)nums.size(); count > 0 && (int)out.size() < k; count--) {
        for (int value : buckets[count]) {
            out.push_back(value);
            if ((int)out.size() == k) return out;
        }
    }
    return out;   // O(n)
}

function topKFrequentBuckets(nums, k) {
  const counts = new Map();
  for (const x of nums) counts.set(x, (counts.get(x) ?? 0) + 1);

  // A frequency can never exceed n, so the bucket array is bounded.
  const buckets = Array.from({ length: nums.length + 1 }, () => []);
  for (const [value, count] of counts) buckets[count].push(value);

  const out = [];
  for (let count = nums.length; count > 0 && out.length < k; count--) {
    for (const value of buckets[count]) {
      out.push(value);
      if (out.length === k) return out;
    }
  }
  return out;   // O(n) — no log factor at all
}

O(n) time. Whenever the key you are sorting by is a bounded integer, bucketing beats a heap — that generalisation is worth more than the specific problem.

Quickselect: O(n) average

Kth Largest Element in an Array is the problem where the interviewer wants to hear

"quickselect". Partition like quicksort, but recurse into only the side containing the answer:

import random

def quickselect(nums, k):
    """kth largest = index len(nums) - k in sorted order."""
    target = len(nums) - k
    lo, hi = 0, len(nums) - 1

    while True:
        pivot_index = random.randint(lo, hi)
        nums[pivot_index], nums[hi] = nums[hi], nums[pivot_index]
        store = lo
        for i in range(lo, hi):
            if nums[i] < nums[hi]:
                nums[i], nums[store] = nums[store], nums[i]
                store += 1
        nums[store], nums[hi] = nums[hi], nums[store]

        if store == target:
            return nums[store]
        if store < target:
            lo = store + 1
        else:
            hi = store - 1

int quickselect(int[] nums, int k) {
    int target = nums.length - k;   // kth largest = this index when sorted
    int lo = 0, hi = nums.length - 1;
    Random rng = new Random();

    while (true) {
        // The random pivot is NOT optional — a fixed one is O(n²) on
        // sorted input, and that is the follow-up question.
        swap(nums, lo + rng.nextInt(hi - lo + 1), hi);

        int store = lo;
        for (int i = lo; i < hi; i++) {
            if (nums[i] < nums[hi]) swap(nums, i, store++);
        }
        swap(nums, store, hi);

        if (store == target) return nums[store];
        if (store < target) lo = store + 1;
        else hi = store - 1;
    }
}

private void swap(int[] a, int i, int j) {
    int t = a[i]; a[i] = a[j]; a[j] = t;
}

int quickselect(vector<int>& nums, int k) {
    int target = (int)nums.size() - k;   // kth largest = this sorted index
    int lo = 0, hi = (int)nums.size() - 1;

    while (true) {
        // A random pivot is required: a fixed one is O(n²) on sorted input.
        swap(nums[lo + rand() % (hi - lo + 1)], nums[hi]);

        int store = lo;
        for (int i = lo; i < hi; i++) {
            if (nums[i] < nums[hi]) swap(nums[i], nums[store++]);
        }
        swap(nums[store], nums[hi]);

        if (store == target) return nums[store];
        if (store < target) lo = store + 1;
        else hi = store - 1;
    }
}

int quickselect(int* nums, int n, int k) {
    int target = n - k;                  /* kth largest = this sorted index */
    int lo = 0, hi = n - 1;

    while (1) {
        int p = lo + rand() % (hi - lo + 1);   /* random pivot is required */
        int tmp = nums[p]; nums[p] = nums[hi]; nums[hi] = tmp;

        int store = lo;
        for (int i = lo; i < hi; i++) {
            if (nums[i] < nums[hi]) {
                tmp = nums[i]; nums[i] = nums[store]; nums[store] = tmp;
                store++;
            }
        }
        tmp = nums[store]; nums[store] = nums[hi]; nums[hi] = tmp;

        if (store == target) return nums[store];
        if (store < target) lo = store + 1;
        else hi = store - 1;
    }
}

function quickselect(nums, k) {
  const target = nums.length - k;   // kth largest = this index when sorted
  let lo = 0, hi = nums.length - 1;
  const swap = (i, j) => { [nums[i], nums[j]] = [nums[j], nums[i]]; };

  while (true) {
    // The random pivot is NOT optional — a fixed one is O(n²) on sorted input.
    swap(lo + Math.floor(Math.random() * (hi - lo + 1)), hi);

    let store = lo;
    for (let i = lo; i < hi; i++) {
      if (nums[i] < nums[hi]) swap(i, store++);
    }
    swap(store, hi);

    if (store === target) return nums[store];
    if (store < target) lo = store + 1;
    else hi = store - 1;
  }
}

Recursing into one side gives n + n/2 + n/4 + … = 2n, so O(n) average. The random pivot is not optional — without it, sorted input degrades to O(n²), and that is the follow-up question.

Picking the approach

Situation	Use	Cost
Streaming, or k ≪ n	Size-K heap	O(n log k)
One-shot, array in memory	Quickselect	O(n) average
Sorting key is a bounded integer	Bucket sort	O(n)
K close to n, or k unknown	Just sort	O(n log n)

K Closest Points to Origin is the heap with squared distance as the key — do not take

the square root, it costs time and changes no ordering. Task Scheduler is a max-heap on remaining counts, or the closed-form idle-slot formula if you spot it.

Complexity

Size-K heap: O(n log k) time, O(k) space. Quickselect: O(n) average, O(n²) worst, O(1) space, but it mutates the input. Bucket sort: O(n) both. Being able to name all three and justify the choice is what the question is actually testing.

Mistakes that cost the round

Using a max-heap for the K largest. The root must be the item you are most willing

to evict.

Letting the heap grow to n. heappush then heappop at size k + 1, not at the

end — otherwise it is O(n log n) with extra steps.

Forgetting to re-negate after the negation trick for a max-heap.
Quickselect with a fixed pivot, which is O(n²) on sorted input.
Taking square roots in K Closest Points.

What to drill

Kth Largest Element in an Array — heap, then quickselect.
Top K Frequent Elements — heap, then bucket sort.
K Closest Points to Origin — the heap with a custom key.
Task Scheduler — max-heap scheduling, or the formula.

All on the 22 DSA Patterns sheet. The two heaps pattern is the next step up.

Frequently asked

Why use a min-heap to find the K largest elements?

Because the root of a size-K min-heap is the smallest element you are currently keeping — the one to evict when something better arrives. Each new element is compared against that single value in O(log k). A max-heap would surface the best element instead, which tells you nothing about what to discard.

Is quickselect better than a heap for top K?

For a one-shot query on an in-memory array, yes: O(n) average versus O(n log k). But it mutates the input, degrades to O(n²) without a random pivot, and cannot handle a stream. A size-K heap is the right answer when data arrives incrementally or k is much smaller than n.

How do you find top K frequent elements in O(n)?

Bucket sort. Count frequencies, then index an array of lists by frequency — no count can exceed n, so the array is bounded. Walk it from the highest count downward and collect until you have K. That avoids the log factor a heap or a sort would add.

Top 'K' Elements