Heap-based Priority Queues and Top-K Problems

1. What is a Priority Queue?

A standard queue follows a First-In, First-Out (FIFO) principle. A Priority Queue (PQ) is different: when you dequeue an element, you don't get the one that has been waiting the longest, but the one with the highest (or lowest) priority.

The main operations are:

• add(element) or enqueue(element): Adds an element to the queue.
• remove() or dequeue(): Removes and returns the element with the highest priority.
• peek(): Returns the highest priority element without removing it.

Think about a hospital emergency room: patients are treated based on the severity of their condition (priority), not just their arrival time. This is a perfect analogy for a priority queue.

So, how can we build one?

• Unsorted Array/List: add is fast ($O(1)$), but finding the highest priority element to remove requires a full scan ($O(N)$).
• Sorted Array/List: remove is fast ($O(1)$), but add requires finding the correct position and shifting elements ($O(N)$).

Neither is ideal. We need a structure that balances both operations. This is where the heap shines. A binary heap is the perfect underlying data structure for a priority queue, offering $O(\log N)$ time for both insertion and deletion.

• A Min-Heap naturally implements a Min-Priority Queue (where remove returns the smallest element).
• A Max-Heap naturally implements a Max-Priority Queue (where remove returns the largest element).

The video below gives an excellent introduction to why heaps are the right choice for this job.

2. Priority Queue Operations via Heap Operations

The core operations of a priority queue map directly to the heap operations you already know.

• add(element): This is equivalent to inserting an element into the heap. We add the element to the end of our array (maintaining the complete binary tree shape) and then perform an upheap (also called sift-up or swim-up) operation to restore the heap property.
• remove(): This is equivalent to removing the root of the heap. We swap the root with the last element, remove the old root (which is now at the end), and then perform a downheap (also called sift-down or heapify) on the new root to restore the heap property.
• peek(): This is a simple $O(1)$ operation that just returns the element at the root of the heap (index 0 of the array).

The following video segments and image illustrate these fundamental mechanics.

3. Using Java's PriorityQueue

For interviews and competitive programming, you'll rarely implement a priority queue from scratch. Instead, you'll use Java's built-in java.util.PriorityQueue. It's crucial to know how to use it effectively.

• By default, PriorityQueue is a Min-Heap. This means poll() will always return the smallest element.
• To create a Max-Heap, you need to provide a custom Comparator. The easiest way is using Collections.reverseOrder() or a lambda expression (a, b) -> b - a.

The article below provides excellent, concise code snippets for common PriorityQueue use cases.

Here's a quick summary of the code:

// Min-heap (default behavior)
PriorityQueue<Integer> minHeap = new PriorityQueue<>();

// Max-heap (using a lambda comparator)
PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a);

// Max-heap (using Collections.reverseOrder())
PriorityQueue<Integer> maxHeap2 = new PriorityQueue<>(Collections.reverseOrder());

// Key methods:
minHeap.offer(10); // Add an element
minHeap.add(5);    // Same as offer

int smallest = minHeap.poll(); // Remove and return the head (5)
int nextSmallest = minHeap.peek(); // Get the head without removing (10)

4. The "Top K Elements" Pattern

Now for the main application. A very common class of interview problems asks you to find the "Top K" largest, smallest, or most frequent items in a large dataset.

The naive approach is to sort the entire dataset ($O(N\log N)$) and pick the first or last K elements. A priority queue offers a much better solution with $O(N\log K)$ time complexity.

The logic is clever and a bit counter-intuitive at first:

• To find the K largest elements, use a Min-Heap of size K.
• To find the K smallest elements, use a Max-Heap of size K.

Let's see why this works for finding the K largest elements:

1. Initialize a min-heap of size K.
2. Iterate through all N elements of the input array.
3. For each element, offer it to the min-heap.
4. If the heap's size grows larger than K, poll() an element. Since it's a min-heap, you are removing the smallest element from the heap.
5. After iterating through all N elements, the heap contains exactly the K largest elements from the entire array.

This process acts as a filter. The heap maintains the "top K candidates" seen so far. When a new, larger element arrives, it's added, and the smallest of the candidates is kicked out, ensuring the K largest are always retained.

This short video explains the pattern beautifully.

5. Applying the Pattern: Solved Problems

Let's put this theory into practice with some classic interview problems.

Problem 1: Top K Frequent Elements (LeetCode 347)

Given an array of integers nums and an integer k, return the k most frequent elements.

This is a two-step problem:

1. Count Frequencies: Use a HashMap to count the occurrences of each number. This takes $O(N)$ time.
2. Find Top K Frequencies: Use the "Top K" pattern on the frequencies. We want the k elements with the largest frequencies, so we will use a min-heap of size k. The heap will store pairs of (number, frequency), ordered by frequency.

Here is what the Java implementation looks like. Notice the custom comparator for the PriorityQueue.

public int[] topKFrequent(int[] nums, int k) {
    // 1. Count frequencies
    Map<Integer, Integer> freqMap = new HashMap<>();
    for (int num : nums) {
        freqMap.put(num, freqMap.getOrDefault(num, 0) + 1);
    }

    // 2. Use a min-heap to find top k frequent elements
    // The heap stores Map.Entry objects, compared by their value (frequency)
    PriorityQueue<Map.Entry<Integer, Integer>> minHeap = 
        new PriorityQueue<>(Comparator.comparingInt(Map.Entry::getValue));

    for (Map.Entry<Integer, Integer> entry : freqMap.entrySet()) {
        minHeap.offer(entry);
        if (minHeap.size() > k) {
            minHeap.poll(); // Remove the entry with the lowest frequency
        }
    }

    // 3. Extract the result from the heap
    int[] result = new int[k];
    for (int i = k - 1; i >= 0; i--) {
        result[i] = minHeap.poll().getKey();
    }
    
    return result;
}

Problem 2: K Closest Points to Origin (LeetCode 973)

Given an array of points and an integer k, find the k closest points to the origin (0, 0).

Here, "priority" is the distance from the origin. We want the k points with the smallest distance. Following our pattern, we should use a max-heap of size k.

1. Initialize a max-heap. The comparator will compare points based on their distance from the origin. To avoid floating-point precision issues and sqrt calculations, we can compare squared distances (x^2 + y^2).
2. Iterate through the points. For each point, offer it to the max-heap.
3. If the heap size exceeds k, poll() the point with the largest distance, which is the root of the max-heap.
4. After the loop, the heap contains the k closest points.

Conclusion

Excellent work! You've not only grasped the concept of a Priority Queue but also learned how to apply it to a powerful and common interview pattern.

Key Takeaways:

• A Priority Queue is an abstract data type that serves elements based on priority, not arrival order.
• The Binary Heap is the ideal data structure for implementing a Priority Queue, providing $O(\log N)$ for insertions and deletions.
• Java's java.util.PriorityQueue is a min-heap by default but can be converted to a max-heap with a custom Comparator.
• The "Top K Elements" pattern efficiently solves problems needing the K largest/smallest items in $O(N\log K)$ time.
  • To find the K largest items, use a min-heap of size K.
  • To find the K smallest items, use a max-heap of size K.

Preview of the Next Lesson:
We will now shift our focus to another incredibly important data structure for interview preparation: the Hash Table. In the next lesson, we will explore the fundamentals of hashing, including hash functions and collision resolution strategies like chaining and open addressing. You'll see how hashing enables constant-time average performance for lookups, insertions, and deletions.