Heap and Priority Queue: Solving Top-K Problems in Java

1// Min-heap (default)
2PriorityQueue<Integer> minHeap = new PriorityQueue<>();
3
4// Max-heap
5PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
6// or: new PriorityQueue<>((a, b) -> b - a)
7
8// Custom comparator — sort by frequency descending
9PriorityQueue<int[]> freqHeap = new PriorityQueue<>((a, b) -> b[1] - a[1]);

1// Approach 1: Min-heap of size K — O(n log k)
2// The top of a min-heap of size K is the Kth largest
3public int findKthLargest(int[] nums, int k) {
4    PriorityQueue<Integer> minHeap = new PriorityQueue<>();
5    for (int num : nums) {
6        minHeap.offer(num);
7        if (minHeap.size() > k) minHeap.poll(); // remove smallest
8    }
9    return minHeap.peek(); // Kth largest
10}
11
12// Approach 2: QuickSelect — O(n) average, O(n²) worst
13// Use when k is small relative to n and you can't afford extra space

1public ListNode mergeKLists(ListNode[] lists) {
2    // Min-heap ordered by node value
3    PriorityQueue<ListNode> heap = new PriorityQueue<>(
4        (a, b) -> a.val - b.val
5    );
6
7    // Add the head of each list
8    for (ListNode node : lists)
9        if (node != null) heap.offer(node);
10
11    ListNode dummy = new ListNode(0), curr = dummy;
12    while (!heap.isEmpty()) {
13        ListNode node = heap.poll();
14        curr.next = node;
15        curr = curr.next;
16        if (node.next != null) heap.offer(node.next);
17    }
18    return dummy.next;
19}

1public int[] topKFrequent(int[] nums, int k) {
2    Map<Integer, Integer> freq = new HashMap<>();
3    for (int n : nums) freq.merge(n, 1, Integer::sum);
4
5    // Min-heap of size k, ordered by frequency
6    PriorityQueue<Map.Entry<Integer, Integer>> heap =
7        new PriorityQueue<>(Comparator.comparingInt(Map.Entry::getValue));
8
9    for (Map.Entry<Integer, Integer> entry : freq.entrySet()) {
10        heap.offer(entry);
11        if (heap.size() > k) heap.poll(); // remove least frequent
12    }
13
14    return heap.stream().mapToInt(Map.Entry::getKey).toArray();
15}
16
17// O(n) bucket sort approach — when k is close to n
18public int[] topKFrequentBucket(int[] nums, int k) {
19    Map<Integer, Integer> freq = new HashMap<>();
20    for (int n : nums) freq.merge(n, 1, Integer::sum);
21
22    List<Integer>[] buckets = new List[nums.length + 1];
23    freq.forEach((num, f) -> {
24        if (buckets[f] == null) buckets[f] = new ArrayList<>();
25        buckets[f].add(num);
26    });
27
28    List<Integer> result = new ArrayList<>();
29    for (int i = buckets.length - 1; i >= 0 && result.size() < k; i--)
30        if (buckets[i] != null) result.addAll(buckets[i]);
31
32    return result.stream().mapToInt(Integer::intValue).toArray();
33}

1class MedianFinder {
2    // Two heaps: maxHeap (lower half) + minHeap (upper half)
3    PriorityQueue<Integer> lo = new PriorityQueue<>(Collections.reverseOrder()); // max-heap
4    PriorityQueue<Integer> hi = new PriorityQueue<>(); // min-heap
5
6    public void addNum(int num) {
7        lo.offer(num);
8        hi.offer(lo.poll()); // balance: push largest of lo into hi
9        if (lo.size() < hi.size()) lo.offer(hi.poll()); // lo always >= hi in size
10    }
11
12    public double findMedian() {
13        return lo.size() > hi.size()
14            ? lo.peek()
15            : (lo.peek() + hi.peek()) / 2.0;
16    }
17}