Learn DSA Part 10: Heaps and Priority Queue - Top K, Merge K Lists and Running Median

Whenever a problem asks for "the k largest", "the smallest so far", "merge many sorted things" or "always process the most urgent item next", the answer is usually a heap. Heaps give you the minimum or maximum in O(1) and let you add or remove items in O(log n), which makes them ideal for streaming data and scheduling.

This is Part 10 of the Learn DSA series. We cover how a heap is stored in an array, how to use priority queues in C++ and Java, and the four heap patterns behind most interview questions: top K, merging K sorted lists, two heaps, and greedy scheduling. Code is in C++ and Java.

1. How a heap works

A binary heap is a complete binary tree (every level full except possibly the last, filled from the left) with the heap property: in a min-heap every parent is less than or equal to its children, so the minimum is always at the root. A max-heap is the mirror image.

Because the tree is complete, it is stored in a plain array with no pointers. For the node at index i: parent is (i - 1) / 2, left child is 2i + 1 and right child is 2i + 2.

• Push: add at the end of the array, then sift up, swapping with the parent while it is smaller.
• Pop: move the last element to the root, then sift down, swapping with the smaller child while needed.

Building a heap bottom-up is O(n), not O(n log n), because most nodes are near the bottom and sift down only a short distance. Heap sort builds a heap and pops n times, giving O(n log n) with O(1) extra space.

2. Priority queues in C++ and Java

For pairs or objects, supply a comparator, for example ordering by distance or frequency. In Java, prefer Integer.compare(a, b) over a - b in comparators, because subtraction can overflow for large values.

3. Pattern 1: top K elements

To find the k largest elements, keep a min-heap of size k. Push each element; if the heap grows beyond k, pop the smallest. At the end the heap holds the k largest and its top is the k-th largest. This takes O(n log k) time and O(k) space, much better than sorting everything when k is small.

// K-th largest element using a min-heap of size k. O(n log k)
int findKthLargest(const vector<int>& nums, int k) {
    priority_queue<int, vector<int>, greater<int>> minHeap;
    for (int x : nums) {
        minHeap.push(x);
        if ((int)minHeap.size() > k) minHeap.pop();   // drop the smallest
    }
    return minHeap.top();
}

// K-th largest element using a min-heap of size k. O(n log k)
static int findKthLargest(int[] nums, int k) {
    PriorityQueue<Integer> minHeap = new PriorityQueue<>();
    for (int x : nums) {
        minHeap.offer(x);
        if (minHeap.size() > k) minHeap.poll();      // drop the smallest
    }
    return minHeap.peek();
}

The same idea finds the K Closest Points to Origin (max-heap of size k by distance) and Top K Frequent Elements (count with a hash map from Part 3, then keep the k most frequent in a heap).

4. Pattern 2: merging K sorted lists

To merge k sorted lists, put the first node of each list in a min-heap. Repeatedly pop the smallest node, append it to the result and push that node's successor. The heap never holds more than k items, so merging N total elements takes O(N log k).

// Merge k sorted linked lists (ListNode from Part 5). O(N log k)
ListNode* mergeKLists(vector<ListNode*>& lists) {
    auto cmp = [](ListNode* a, ListNode* b) { return a->val > b->val; };
    priority_queue<ListNode*, vector<ListNode*>, decltype(cmp)> pq(cmp);
    for (ListNode* head : lists) if (head) pq.push(head);
    ListNode dummy(0);
    ListNode* tail = &dummy;
    while (!pq.empty()) {
        ListNode* node = pq.top(); pq.pop();
        tail->next = node;
        tail = node;
        if (node->next) pq.push(node->next);
    }
    return dummy.next;
}

// Merge k sorted linked lists (ListNode from Part 5). O(N log k)
static ListNode mergeKLists(ListNode[] lists) {
    PriorityQueue<ListNode> pq =
        new PriorityQueue<>((a, b) -> Integer.compare(a.val, b.val));
    for (ListNode head : lists) if (head != null) pq.offer(head);
    ListNode dummy = new ListNode(0), tail = dummy;
    while (!pq.isEmpty()) {
        ListNode node = pq.poll();
        tail.next = node;
        tail = node;
        if (node.next != null) pq.offer(node.next);
    }
    return dummy.next;
}

The same pattern finds the Kth Smallest Element in a Sorted Matrix and the smallest range covering elements from k lists.

5. Pattern 3: two heaps for the running median

To report the median of a stream as numbers arrive, keep the smaller half in a max-heap and the larger half in a min-heap, with sizes differing by at most one. The median is the top of the bigger heap, or the average of both tops. Each insertion is O(log n) and each median query O(1).

// Median of a data stream with two heaps
class MedianFinder {
    priority_queue<int> low;                                  // max-heap
    priority_queue<int, vector<int>, greater<int>> high;      // min-heap
public:
    void addNum(int x) {
        low.push(x);
        high.push(low.top()); low.pop();           // keep order between halves
        if (high.size() > low.size()) { low.push(high.top()); high.pop(); }
    }
    double findMedian() {
        if (low.size() > high.size()) return low.top();
        return (low.top() + (double)high.top()) / 2.0;
    }
};

// Median of a data stream with two heaps
class MedianFinder {
    PriorityQueue<Integer> low = new PriorityQueue<>(Collections.reverseOrder());
    PriorityQueue<Integer> high = new PriorityQueue<>();
    void addNum(int x) {
        low.offer(x);
        high.offer(low.poll());                    // keep order between halves
        if (high.size() > low.size()) low.offer(high.poll());
    }
    double findMedian() {
        if (low.size() > high.size()) return low.peek();
        return (low.peek() + (double) high.peek()) / 2.0;
    }
}

Sliding Window Median and IPO-style "pick the best affordable project" problems use the same two-heap balancing.

6. Pattern 4: greedy scheduling with a heap

• Meeting Rooms II: sort meetings by start time and keep a min-heap of end times. If the earliest-ending room is free, reuse it; otherwise add a room. The heap size is the answer.
• Task Scheduler / Reorganize String: always pick the most frequent remaining item next, using a max-heap of counts.
• Dijkstra's shortest path: repeatedly expand the closest unvisited node, using a min-heap of distances (Part 11).

Practice problems

These are classic LeetCode-style problems, listed from easy to hard. For each one, decide whether you need a min-heap or a max-heap and what it should store.

1. Last Stone Weight (easy)
2. Kth Largest Element in a Stream (easy)
3. Kth Largest Element in an Array (medium)
4. K Closest Points to Origin (medium)
5. Top K Frequent Elements (medium)
6. Task Scheduler (medium)
7. Reorganize String (medium)
8. Meeting Rooms II (medium)
9. Kth Smallest Element in a Sorted Matrix (medium)
10. Merge k Sorted Lists (hard)
11. Find Median from Data Stream (hard)
12. Sliding Window Median (hard)
13. IPO (hard)