Advanced Heap Patterns: Merge K Lists, Median, Top K

---

Pattern 1: Merge K Sorted Lists

The Problem

Given K sorted lists (or arrays), merge them into one sorted output.

Why a Heap?

The naive approach merges lists pairwise: O(N * K). A min-heap approach only keeps K elements in the heap at any time, extracting the smallest and pushing the next element from the same list.

Time: O(N log K) where N is the total number of elements.

Implementation

import heapq

def merge_k_sorted_lists(lists):
    """
    Merge K sorted lists into one sorted list.
    Time: O(N log K), Space: O(K) for the heap.
    """
    result = []
    # Min-heap of (value, list_index, element_index)
    heap = []
    
    # Initialize heap with the first element of each list
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))
    
    while heap:
        val, list_idx, elem_idx = heapq.heappop(heap)
        result.append(val)
        
        # Push the next element from the same list
        next_idx = elem_idx + 1
        if next_idx < len(lists[list_idx]):
            heapq.heappush(heap, (lists[list_idx][next_idx], list_idx, next_idx))
    
    return result


lists = [
    [1, 4, 7, 10],
    [2, 5, 8, 11],
    [3, 6, 9, 12],
]
print(merge_k_sorted_lists(lists))
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

Linked List Version (LeetCode 23)

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

def merge_k_linked_lists(lists):
    """Merge K sorted linked lists. LeetCode 23."""
    heap = []
    
    # Use a counter to break ties (ListNode is not comparable)
    counter = 0
    for head in lists:
        if head:
            heapq.heappush(heap, (head.val, counter, head))
            counter += 1
    
    dummy = ListNode(0)
    current = dummy
    
    while heap:
        val, _, node = heapq.heappop(heap)
        current.next = node
        current = current.next
        
        if node.next:
            counter += 1
            heapq.heappush(heap, (node.next.val, counter, node.next))
    
    return dummy.next

---

Pattern 2: Find Median from Data Stream (Two Heaps)

The Problem

Design a data structure that supports:

• add_num(num): Add a number from the stream.
• find_median(): Return the median of all numbers seen so far.

The Two-Heap Approach

Maintain two heaps:

• max_heap (stores the smaller half, negated for Python’s min-heap)
• min_heap (stores the larger half)

Invariant: len(max_heap) == len(min_heap) or len(max_heap) == len(min_heap) + 1

The median is either the top of max_heap (odd total) or the average of both tops (even total).

class MedianFinder:
    """Find median from data stream using two heaps. LeetCode 295."""
    
    def __init__(self):
        self.max_heap = []  # Stores negated values (smaller half)
        self.min_heap = []  # Stores values (larger half)
    
    def add_num(self, num):
        """Add a number. O(log n)."""
        # Always add to max_heap first
        heapq.heappush(self.max_heap, -num)
        
        # Ensure max_heap's top <= min_heap's top
        if self.min_heap and (-self.max_heap[0]) > self.min_heap[0]:
            val = -heapq.heappop(self.max_heap)
            heapq.heappush(self.min_heap, val)
        
        # Balance sizes: max_heap can have at most 1 more element
        if len(self.max_heap) > len(self.min_heap) + 1:
            val = -heapq.heappop(self.max_heap)
            heapq.heappush(self.min_heap, val)
        elif len(self.min_heap) > len(self.max_heap):
            val = heapq.heappop(self.min_heap)
            heapq.heappush(self.max_heap, -val)
    
    def find_median(self):
        """Return current median. O(1)."""
        if len(self.max_heap) > len(self.min_heap):
            return -self.max_heap[0]
        return (-self.max_heap[0] + self.min_heap[0]) / 2.0


# Example
mf = MedianFinder()
for num in [3, 1, 5, 4, 2]:
    mf.add_num(num)
    print(f"After adding {num}: median = {mf.find_median()}")

Output:

After adding 3: median = 3
After adding 1: median = 2.0
After adding 5: median = 3
After adding 4: median = 3.5
After adding 2: median = 3

---

Pattern 3: Top K Frequent Elements

The Problem

Given an array, return the K most frequent elements.

Approach: Count + Min-Heap of Size K

from collections import Counter

def top_k_frequent(nums, k):
    """
    Return the K most frequent elements. LeetCode 347.
    Time: O(n log k), Space: O(n) for counter.
    """
    count = Counter(nums)
    
    # Min-heap of size K: (frequency, element)
    # We keep the K largest frequencies
    heap = []
    
    for num, freq in count.items():
        heapq.heappush(heap, (freq, num))
        if len(heap) > k:
            heapq.heappop(heap)  # Remove smallest frequency
    
    return [num for freq, num in heap]


nums = [1, 1, 1, 2, 2, 3, 3, 3, 3]
print(top_k_frequent(nums, 2))  # [1, 3]

Alternative: Bucket Sort for O(n)

def top_k_frequent_bucket(nums, k):
    """O(n) solution using bucket sort."""
    count = Counter(nums)
    
    # Bucket: index = frequency, value = list of elements with that frequency
    buckets = [[] for _ in range(len(nums) + 1)]
    for num, freq in count.items():
        buckets[freq].append(num)
    
    result = []
    for freq in range(len(buckets) - 1, 0, -1):
        for num in buckets[freq]:
            result.append(num)
            if len(result) == k:
                return result
    
    return result

---

Pattern 4: K Closest Points to Origin

def k_closest(points, k):
    """
    Return K closest points to origin. LeetCode 973.
    Use max-heap of size K to keep the K smallest distances.
    Time: O(n log k).
    """
    # Max-heap (negate distance): keep K smallest
    heap = []
    
    for x, y in points:
        dist = -(x * x + y * y)  # Negate for max-heap
        
        if len(heap) < k:
            heapq.heappush(heap, (dist, x, y))
        elif dist > heap[0][0]:  # Closer than farthest in heap
            heapq.heapreplace(heap, (dist, x, y))
    
    return [[x, y] for _, x, y in heap]


points = [[3, 3], [5, -1], [-2, 4], [1, 1], [0, 2]]
print(k_closest(points, 3))  # [[1, 1], [0, 2], [-2, 4]] or similar

---

Pattern 5: Reorganize String

Given a string, rearrange so no two adjacent characters are the same. Return empty string if impossible.

def reorganize_string(s):
    """
    Reorganize string so no adjacent chars are same. LeetCode 767.
    Greedy: always pick the most frequent character that isn't the previous one.
    Time: O(n log k) where k = distinct characters.
    """
    count = Counter(s)
    
    # Check if possible: no character can appear more than (n+1)/2 times
    max_freq = max(count.values())
    if max_freq > (len(s) + 1) // 2:
        return ""
    
    # Max-heap of (-frequency, character)
    heap = [(-freq, ch) for ch, freq in count.items()]
    heapq.heapify(heap)
    
    result = []
    prev_freq, prev_ch = 0, ''
    
    while heap:
        freq, ch = heapq.heappop(heap)
        
        # Push back previous character if it still has count
        if prev_freq < 0:
            heapq.heappush(heap, (prev_freq, prev_ch))
        
        result.append(ch)
        prev_freq = freq + 1  # Used one occurrence (freq is negative)
        prev_ch = ch
    
    return ''.join(result)


print(reorganize_string("aab"))     # "aba"
print(reorganize_string("aaab"))    # ""
print(reorganize_string("aaabbc"))  # "ababac" or similar

---

Pattern 6: Task Scheduler

Given tasks and a cooldown period n, find the minimum intervals to finish all tasks.

def least_interval(tasks, n):
    """
    Task scheduler with cooldown. LeetCode 621.
    Time: O(total_tasks * n) in worst case.
    """
    count = Counter(tasks)
    
    # Max-heap of remaining counts
    heap = [-freq for freq in count.values()]
    heapq.heapify(heap)
    
    time = 0
    cooldown_queue = []  # (available_time, remaining_count)
    
    while heap or cooldown_queue:
        time += 1
        
        if heap:
            remaining = heapq.heappop(heap) + 1  # Execute one task
            if remaining < 0:
                cooldown_queue.append((time + n, remaining))
        
        # Check if any task has finished its cooldown
        if cooldown_queue and cooldown_queue[0][0] == time:
            _, remaining = cooldown_queue.pop(0)
            heapq.heappush(heap, remaining)
    
    return time


# Also: math formula approach (more efficient)
def least_interval_math(tasks, n):
    """O(n) solution using math."""
    count = Counter(tasks)
    max_freq = max(count.values())
    max_count = sum(1 for freq in count.values() if freq == max_freq)
    
    # Formula: (max_freq - 1) * (n + 1) + max_count
    # But answer is at least len(tasks)
    return max(len(tasks), (max_freq - 1) * (n + 1) + max_count)


tasks = ['A', 'A', 'A', 'B', 'B', 'B']
print(least_interval(tasks, 2))       # 8
print(least_interval_math(tasks, 2))  # 8

---

Pattern 7: Kth Largest Element in a Stream

class KthLargest:
    """
    Kth largest element in a stream. LeetCode 703.
    Maintain a min-heap of size K.
    The root is always the Kth largest.
    """
    
    def __init__(self, k, nums):
        self.k = k
        self.heap = []
        for num in nums:
            self.add(num)
    
    def add(self, val):
        """Add val and return current Kth largest. O(log k)."""
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]


kl = KthLargest(3, [4, 5, 8, 2])
print(kl.add(3))   # 4
print(kl.add(5))   # 5
print(kl.add(10))  # 5
print(kl.add(9))   # 8
print(kl.add(4))   # 8

---

Pattern 8: Merge K Sorted Arrays with Iterator

When arrays are too large to fit in memory, process them as iterators:

def merge_k_iterators(iterators):
    """
    Lazily merge K sorted iterators.
    Yields elements one at a time without loading everything into memory.
    """
    heap = []
    
    for i, it in enumerate(iterators):
        try:
            val = next(it)
            heapq.heappush(heap, (val, i, it))
        except StopIteration:
            pass
    
    while heap:
        val, idx, it = heapq.heappop(heap)
        yield val
        
        try:
            next_val = next(it)
            heapq.heappush(heap, (next_val, idx, it))
        except StopIteration:
            pass


# Usage with generators
def gen_range(start, end, step=1):
    i = start
    while i < end:
        yield i
        i += step

iters = [
    gen_range(0, 10, 3),   # 0, 3, 6, 9
    gen_range(1, 10, 3),   # 1, 4, 7
    gen_range(2, 10, 3),   # 2, 5, 8
]

print(list(merge_k_iterators(iters)))
# [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

---

When to Use Heap vs Sorting

Scenario	Heap	Sort
Need top K of N elements	O(N log K)	O(N log N)
Data arrives as stream	Yes (online)	No (need all data)
Need all elements sorted	O(N log N)	O(N log N) — prefer sort
Frequently adding/removing	O(log N) per op	O(N) to re-sort
Memory-constrained (K << N)	O(K) heap	O(N) for sorted array

Rule of thumb: Use a heap when K is much smaller than N, or when data arrives incrementally. Use sorting when you need the full sorted output and have all data up front.

---

heapq Cheat Sheet

import heapq

# Create heap from list
nums = [5, 3, 8, 1, 9]
heapq.heapify(nums)         # In-place, O(n)

# Push and pop
heapq.heappush(nums, 2)     # O(log n)
smallest = heapq.heappop(nums)  # O(log n)

# Push then pop (more efficient than separate operations)
result = heapq.heappushpop(nums, 4)  # Push 4, pop smallest

# Pop then push
result = heapq.heapreplace(nums, 7)  # Pop smallest, push 7

# K smallest/largest (O(n + k log n))
heapq.nsmallest(3, [5, 3, 8, 1, 9])  # [1, 3, 5]
heapq.nlargest(3, [5, 3, 8, 1, 9])   # [9, 8, 5]

# Max-heap trick: negate values
max_heap = []
heapq.heappush(max_heap, -5)
heapq.heappush(max_heap, -3)
largest = -heapq.heappop(max_heap)  # 5

# Custom comparison with tuples
heapq.heappush([], (priority, counter, item))

---

Practice Problems

1. Merge K Sorted Lists (LeetCode 23) — Min-heap with K pointers.
2. Find Median from Data Stream (LeetCode 295) — Two heaps.
3. Top K Frequent Elements (LeetCode 347) — Count + min-heap of size K.
4. K Closest Points to Origin (LeetCode 973) — Max-heap of size K.
5. Reorganize String (LeetCode 767) — Greedy with max-heap.
6. Task Scheduler (LeetCode 621) — Max-heap + cooldown queue.
7. Kth Largest Element in Stream (LeetCode 703) — Min-heap of size K.
8. Sort Characters by Frequency (LeetCode 451) — Count + max-heap.
9. Ugly Number II (LeetCode 264) — Min-heap for generating sequence.
10. Smallest Range Covering Elements from K Lists (LeetCode 632) — Min-heap + sliding window.

---

Key Takeaways

• The merge K sorted lists pattern uses a min-heap of size K to always extract the global minimum. Time: O(N log K).
• The two-heap median pattern maintains a max-heap (smaller half) and min-heap (larger half) for O(log n) insertions and O(1) median queries.
• The top K pattern uses a min-heap of size K: elements smaller than the root are discarded, so only the K largest survive.
• Python’s heapq is a min-heap; simulate max-heap by negating values.
• Heaps shine when K is much smaller than N or when data arrives as a stream.