Top-K Elements

Top-K Elements

"Find the k largest", "the k most frequent words", "the k closest points", "the k-th smallest": whenever a question asks for the best k items out of many, fully sorting the input (O(n log n)) does more work than needed. The top-k pattern keeps only the k candidates that currently qualify, usually in a heap of size k, giving O(n log k). QuickSelect can even find the k-th element in O(n) on average.

The core trick: a heap of the opposite kind

To keep the k largest items, use a min-heap of size k:

stream: 7, 2, 9, 4, 11, 5        k = 3
heap (min at top) after each:
  [7] → [2,7] → [2,7,9] → add 4: [2,4,7,9] → evict min 2 → [4,7,9]
  → add 11 → evict 4 → [7,9,11] → add 5 → evict 5 → [7,9,11]

The heap's top is the smallest of the current top k, so it's the one to evict when something better arrives. At the end, the top is the k-th largest. The mirror rule also holds: for the k smallest, use a max-heap of size k.

int findKthLargest(int[] nums, int k) {
    PriorityQueue<Integer> heap = new PriorityQueue<>();        // Java's default is a MIN-heap
    for (int x : nums) {
        heap.offer(x);
        if (heap.size() > k) heap.poll();                        // drop the smallest → keep the k largest
    }
    return heap.peek();
}

O(n log k) time and O(k) space. When k is much smaller than n, this is much better than sorting. It also works on streams: you never need all the data in memory at once, which matters for log files and event feeds.

QuickSelect: the k-th element in O(n) average

QuickSelect uses quicksort's partition step, but recurses into only the side that contains the answer:

int findKthLargest(int[] nums, int k) {
    int target = nums.length - k;                   // k-th largest = index n−k in sorted order
    int lo = 0, hi = nums.length - 1;
    Random rnd = new Random();
    while (true) {
        int p = partition(nums, lo, hi, lo + rnd.nextInt(hi - lo + 1));   // random pivot
        if (p == target) return nums[p];
        if (p < target) lo = p + 1; else hi = p - 1;
    }
}

private int partition(int[] a, int lo, int hi, int pivotIdx) {
    int pivot = a[pivotIdx];
    swap(a, pivotIdx, hi);
    int store = lo;
    for (int i = lo; i < hi; i++) if (a[i] < pivot) swap(a, i, store++);
    swap(a, store, hi);
    return store;                                   // pivot's final sorted position
}

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

• The average work shrinks geometrically (n + n/2 + n/4 + … ≈ 2n), so it's O(n) on average.
• With consistently bad pivots, for example always the largest, it degrades to O(n²). A random pivot makes that vanishingly unlikely. The median-of-medians algorithm guarantees O(n), but is rarely worth its complexity.
• It rearranges the input array and needs all the data in memory, unlike the heap approach.

Top K Frequent Elements

Count frequencies, then select the k highest counts.

With a heap (O(n log k)):

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

    PriorityQueue<Map.Entry<Integer, Integer>> heap =
            new PriorityQueue<>(Map.Entry.comparingByValue());      // min-heap by frequency
    for (Map.Entry<Integer, Integer> e : count.entrySet()) {
        heap.offer(e);
        if (heap.size() > k) heap.poll();
    }
    return heap.stream().mapToInt(Map.Entry::getKey).toArray();
}

With bucket sort (O(n)): a frequency is at most n, so make an array of buckets indexed by frequency and read it from the highest bucket down:

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

    List<List<Integer>> buckets = new ArrayList<>();
    for (int i = 0; i <= nums.length; i++) buckets.add(new ArrayList<>());
    count.forEach((value, freq) -> buckets.get(freq).add(value));

    int[] result = new int[k];
    int idx = 0;
    for (int f = nums.length; f >= 1 && idx < k; f--)
        for (int v : buckets.get(f)) { if (idx == k) break; result[idx++] = v; }
    return result;
}

Bucket sort works because the key (a frequency) is a bounded integer. It's worth offering as the linear-time follow-up.

Top K Frequent Words adds a tie-break: equal frequencies are ordered alphabetically. The heap comparator must then say "lower frequency first, and for equal frequency, the alphabetically later word first", so the right item is evicted. Getting this comparator exactly right is the whole challenge of that problem.

K Closest Points to Origin

"Closest" means smallest distance, so keep a max-heap of size k by distance and evict the farthest. Compare squared distances: no sqrt is needed, since squaring preserves order. Use long to avoid overflow for large coordinates:

int[][] kClosest(int[][] points, int k) {
    PriorityQueue<int[]> heap = new PriorityQueue<>(
            (a, b) -> Long.compare(dist2(b), dist2(a)));            // max-heap: farthest on top
    for (int[] p : points) {
        heap.offer(p);
        if (heap.size() > k) heap.poll();
    }
    return heap.toArray(new int[0][]);
}
private long dist2(int[] p) { return (long) p[0] * p[0] + (long) p[1] * p[1]; }

Comparator pitfalls

• (a, b) -> a - b overflows for values near Integer.MIN_VALUE/MAX_VALUE and can silently produce the wrong order. Use Integer.compare(a, b), Comparator.comparingInt(...) or Comparator.reverseOrder().
• PriorityQueue's iterator and toArray return elements in heap order, not sorted order. If the output must be sorted, poll the elements out one by one (or sort the k results, which is cheap).

Where it shows up in real systems

Trending topics and "most viewed" lists over streams, often with approximate counting (the Count-Min Sketch plus a heap) when the number of distinct keys is huge. Leaderboards. Search engines returning the top-k scored documents from millions of candidates. Databases executing ORDER BY … LIMIT k with a bounded heap instead of a full sort.

Follow-up questions this topic invites — and their answers

Q: Why use a min-heap to find the largest elements? A: A heap of size k holding the k largest seen so far must evict its smallest member when a larger item arrives. A min-heap keeps that smallest member at the top, so eviction is O(log k), and at the end the top is exactly the k-th largest.

Q: Heap or QuickSelect? A: QuickSelect is O(n) on average and uses O(1) extra space, but it modifies the input, needs all data in memory, and has an O(n²) worst case. The heap is O(n log k), works on streams, and leaves the input untouched. In interviews, give the heap solution first and mention QuickSelect as the average-linear alternative.

Q: How do you make QuickSelect's worst case unlikely? A: Choose the pivot at random (or shuffle the array first), and use a partition scheme that handles many duplicates well (three-way partitioning). Median-of-medians guarantees O(n) but has large constant factors.

Q: What if the data is too large to fit in memory? A: Stream it through a size-k heap, which is O(k) memory. For top-k frequent items over a huge stream with many distinct keys, exact counting may not fit either. Use approximate sketches (Count-Min Sketch, the Space-Saving algorithm), or count per shard in a distributed job and merge the per-shard top-k lists.

Q: Can you get the top k in sorted order efficiently? A: Yes. After the heap pass, poll the k elements, which comes out ascending for a min-heap, and reverse if needed. That adds O(k log k), which is negligible when k is small.