Advanced Heap and Priority Queue: Top-K and Streaming Data

When to Use a Heap (Priority Queue)

• Top K Problems: Maintain a heap of size $K$.
• Median from Data Stream: Dual-Heap system (Max Heap + Min Heap).
• K-way Merge: Merging multiple sorted sequences efficiently.
• Task Scheduling: Processing tasks based on dynamic priority.

Top K Patterns

Kth Largest Element in a Stream (LeetCode 703)

Maintain exactly $k$ elements in a Min Heap. The root of this heap will always be the $k$-th largest element in the entire stream.

class KthLargest {
    private final PriorityQueue<Integer> pq = new PriorityQueue<>(); // Min Heap
    private final int k;

    KthLargest(int k, int[] nums) {
        this.k = k;
        for (int n : nums) add(n);
    }

    int add(int val) {
        pq.offer(val);
        // If we exceed k elements, the smallest is not in the top K
        if (pq.size() > k) {
            pq.poll();
        }
        return pq.peek(); // Top of k largest elements is the k-th largest overall
    }
}

Rule of Thumb: To find the "K Largest" elements, use a Min Heap (it discards lower-value elements). To find the "K Smallest," use a Max Heap.

Find K Pairs with Smallest Sums (LeetCode 373)

K-way merge logic: Initialize by adding all (nums1[i] + nums2[0]) pairs to the heap. Whenever a pair is popped, add the pair with the same nums1[i] but incrementing the nums2 index.

List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
    List<List<Integer>> result = new ArrayList<>();
    // Store [sum, index_in_nums1, index_in_nums2]
    PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[0] - b[0]);

    // Initial population: first element of nums2 paired with each relevant nums1
    for (int i = 0; i < Math.min(nums1.length, k); i++) {
        pq.offer(new int[]{nums1[i] + nums2[0], i, 0});
    }

    while (!pq.isEmpty() && result.size() < k) {
        int[] current = pq.poll();
        int i = current[1], j = current[2];
        result.add(Arrays.asList(nums1[i], nums2[j]));
        
        // Add the next pair in the "sequence" starting with nums1[i]
        if (j + 1 < nums2.length) {
            pq.offer(new int[]{nums1[i] + nums2[j + 1], i, j + 1});
        }
    }
    return result;
}

Median from Data Stream (LeetCode 295)

Maintain a dynamic balance between two heaps:

• Max Heap: Stores the smaller half of numbers.
• Min Heap: Stores the larger half of numbers.

Invariant: MaxHeap.size() $\ge$ MinHeap.size(), with difference $\le 1$.

class MedianFinder {
    // Stores the smaller half (Max Heap)
    PriorityQueue<Integer> left = new PriorityQueue<>(Collections.reverseOrder());
    // Stores the larger half (Min Heap)
    PriorityQueue<Integer> right = new PriorityQueue<>();

    void addNum(int num) {
        left.offer(num);
        // Move the current largest of 'left' to 'right' to ensure sorting
        right.offer(left.poll());
        
        // Maintain the size invariant
        if (right.size() > left.size()) {
            left.offer(right.poll());
        }
    }

    double findMedian() {
        if (left.size() > right.size()) return left.peek();
        return (left.peek() + right.peek()) / 2.0;
    }
}

Sliding Window Maximum (LeetCode 239)

While a heap can work (O(N log N)), a Monotonic Decreasing Deque is the optimal O(N) solution:

int[] maxSlidingWindow(int[] nums, int k) {
    Deque<Integer> dq = new ArrayDeque<>(); // Stores INDICES
    int[] results = new int[nums.length - k + 1];

    for (int i = 0; i < nums.length; i++) {
        // 1. Remove indices that are outside the current window
        while (!dq.isEmpty() && dq.peekFirst() < i - k + 1) {
            dq.pollFirst();
        }
        // 2. Maintain Monotonic Decreasing property: remove values smaller than current
        while (!dq.isEmpty() && nums[dq.peekLast()] < nums[i]) {
            dq.pollLast();
        }
        dq.offerLast(i);
        
        // 3. Current max is always at the front of the deque
        if (i >= k - 1) {
            results[i - k + 1] = nums[dq.peekFirst()];
        }
    }
    return results;
}

Merge K Sorted Lists (LeetCode 23)

Standard K-way merge using a priority queue for the heads of all lists:

ListNode mergeKLists(ListNode[] lists) {
    PriorityQueue<ListNode> pq = new PriorityQueue<>(Comparator.comparingInt(n -> n.val));
    for (ListNode node : lists) {
        if (node != null) pq.offer(node);
    }

    ListNode dummy = new ListNode(0), current = dummy;
    while (!pq.isEmpty()) {
        ListNode minNode = pq.poll();
        current.next = minNode;
        current = current.next;
        if (minNode.next != null) pq.offer(minNode.next);
    }
    return dummy.next;
}

Selection Summary