Solve 8 classic heap interview problems — kth largest, top K frequent words, merge K lists, reorganize string, IPO (max capital), find median from data stream, smallest range covering K lists, and minimum cost to connect ropes.
Problem 1 — Kth Largest Element in Stream
Problem: Design a class to find the kth largest element in a stream.
import heapq
class KthLargest:
"""
Min-heap of size k: root is the kth largest.
add(): push new element, pop if size > k. O(log k).
"""
def __init__(self, k, nums):
self.k = k
self.heap = []
for num in nums:
self.add(num)
def add(self, val):
heapq.heappush(self.heap, val)
if len(self.heap) > self.k:
heapq.heappop(self.heap)
return self.heap[0] # Root = kth largest
kth = KthLargest(3, [4, 5, 8, 2])
print(kth.add(3)) # 4
print(kth.add(5)) # 5
print(kth.add(10)) # 5
print(kth.add(9)) # 8
print(kth.add(4)) # 8
Problem 3 — Merge K Sorted Lists
import heapq
class ListNode:
def __init__(self, val=0, next=None):
self.val = val
self.next = next
def __lt__(self, other):
return self.val < other.val # For heap comparison
def merge_k_lists(lists):
"""
Merge k sorted linked lists. O(n log k).
Push heads of all lists into min-heap.
Always extract minimum, add its next node.
"""
heap = []
for head in lists:
if head:
heapq.heappush(heap, head)
dummy = ListNode(0)
current = dummy
while heap:
node = heapq.heappop(heap)
current.next = node
current = current.next
if node.next:
heapq.heappush(heap, node.next)
return dummy.next
Problem 4 — Reorganize String
Problem: Rearrange string so no two adjacent characters are the same. Return "" if impossible.
import heapq
from collections import Counter
def reorganize_string(s):
"""
Greedy: always place the most frequent remaining character.
Use max-heap to get most frequent character.
O(n log n) where n = len(s).
"""
freq = Counter(s)
max_heap = [(-count, char) for char, count in freq.items()]
heapq.heapify(max_heap)
result = []
prev_count, prev_char = 0, '' # Track previous character
while max_heap:
count, char = heapq.heappop(max_heap)
if char == prev_char:
# Cannot use most frequent — must use second most frequent
if not max_heap:
return "" # Impossible
count2, char2 = heapq.heappop(max_heap)
result.append(char2)
if count2 + 1 < 0: # Still has remaining instances
heapq.heappush(max_heap, (count2 + 1, char2))
heapq.heappush(max_heap, (count, char)) # Push back original
else:
result.append(char)
prev_count, prev_char = count, char
if count + 1 < 0:
heapq.heappush(max_heap, (count + 1, char))
return ''.join(result)
print(reorganize_string("aab")) # "aba"
print(reorganize_string("aaab")) # ""
Problem 5 — IPO (Maximize Capital)
Problem: Given k projects with profits and capital requirements, start with initial capital w and select at most k projects to maximize final capital.
import heapq
def find_maximized_capital(k, w, profits, capital):
"""
Greedy: always pick the most profitable project we can currently afford.
Min-heap for capital (to find available projects efficiently).
Max-heap for profits (to always pick highest profit available).
O(n log n + k log n)
"""
projects = sorted(zip(capital, profits)) # Sort by capital requirement
available = [] # Max-heap of profits
i = 0
for _ in range(k):
# Add all projects we can now afford
while i < len(projects) and projects[i][0] <= w:
heapq.heappush(available, -projects[i][1]) # Max-heap
i += 1
if not available:
break # Can't afford any more projects
w += -heapq.heappop(available) # Execute most profitable
return w
print(find_maximized_capital(2, 0, [1,2,3], [0,1,1])) # 4
print(find_maximized_capital(3, 0, [1,2,3], [0,1,2])) # 6
Problem 6 — Smallest Range Covering K Lists
Problem: Given k sorted lists, find the smallest range [a,b] such that each list has at least one element in [a,b].
import heapq
def smallest_range(nums):
"""
Use min-heap with one element from each list.
Track current range [min_val, max_val].
Advance the minimum element's list to shrink range.
O(n log k) where n = total elements.
"""
# Initialize: first element from each list
heap = [(row[0], i, 0) for i, row in enumerate(nums)]
heapq.heapify(heap)
current_max = max(row[0] for row in nums)
best = [heap[0][0], current_max]
while True:
current_min, list_idx, elem_idx = heapq.heappop(heap)
# Check if current range is better
if current_max - current_min < best[1] - best[0]:
best = [current_min, current_max]
# Advance to next element in the same list
if elem_idx + 1 == len(nums[list_idx]):
break # This list is exhausted — can't cover all lists
next_val = nums[list_idx][elem_idx + 1]
current_max = max(current_max, next_val)
heapq.heappush(heap, (next_val, list_idx, elem_idx + 1))
return best
print(smallest_range([[4,10,15,24,26],[0,9,12,20],[5,18,22,30]]))
# [20, 24]
Problem 7 — Minimum Cost to Connect Ropes
Problem: Given ropes of different lengths, find minimum cost to connect all. Cost of connecting two ropes = sum of their lengths.
import heapq
def connect_ropes(lengths):
"""
Greedy: always combine two shortest ropes.
Min-heap to get two shortest in O(log n) each time.
O(n log n) total.
"""
heapq.heapify(lengths)
total_cost = 0
while len(lengths) > 1:
first = heapq.heappop(lengths)
second = heapq.heappop(lengths)
cost = first + second
total_cost += cost
heapq.heappush(lengths, cost)
return total_cost
print(connect_ropes([4, 3, 2, 6])) # 29
# (2+3)=5, (4+5)=9, (6+9)=15 → total = 5+9+15 = 29
import heapq
class MedianFinder:
"""
Two heaps approach: O(log n) add, O(1) findMedian.
"""
def __init__(self):
self.low = [] # Max-heap: lower half (negate values)
self.high = [] # Min-heap: upper half
def add_num(self, num):
heapq.heappush(self.low, -num)
# Ensure all low ≤ all high
if self.high and -self.low[0] > self.high[0]:
heapq.heappush(self.high, -heapq.heappop(self.low))
# Balance sizes
if len(self.low) > len(self.high) + 1:
heapq.heappush(self.high, -heapq.heappop(self.low))
elif len(self.high) > len(self.low):
heapq.heappush(self.low, -heapq.heappop(self.high))
def find_median(self):
if len(self.low) == len(self.high):
return (-self.low[0] + self.high[0]) / 2
return float(-self.low[0])
mf = MedianFinder()
for n in [5, 15, 1, 3]:
mf.add_num(n)
print(mf.find_median())
# 5, 10.0, 5, 4.0
Pattern Summary
| Problem | Heap Type | Key Pattern |
|---|
| Kth Largest Stream | Min-heap (size k) | Root = kth largest |
| Top K Frequent Words | Max-heap | Freq-word tuples for ordering |
| Merge K Lists | Min-heap | Always extract current minimum |
| Reorganize String | Max-heap | Greedy: most frequent first |
| IPO | Min-heap + Max-heap | Unlock projects, pick best profit |
| Smallest Range | Min-heap | Advance minimum to shrink range |
| Connect Ropes | Min-heap | Huffman-style greedy |
| Median Stream | Two heaps | Balance lower/upper halves |
Heap problems follow this pattern: Repeatedly identify an extreme (minimum or maximum) of a changing set, process it, and update the set. A heap makes each "find + remove + add" sequence O(log n).
*Heaps section complete. Moving to Trees.*