Median from Data Stream and K-Way Merge
Maintain two heaps (max-heap of small half, min-heap of large half) for O(log n) median updates, and merge k sorted lists using a heap.
Median from Data Stream and K-Way Merge is a free DSA Interview Prep lesson on CoddyKit — lesson 4 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the DSA Interview Prep learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.
Median from Data Stream Problem
Find Median from Data Stream (LeetCode #295) asks you to support two operations efficiently: addNum(num) to add a number and findMedian() to return the current median. The median of an even-length list is the average of the two middle values. A brute-force sorted list gives O(n) insert and O(1) median. The optimal solution uses two heaps for O(log n) insert and O(1) median.
import heapq
# Strategy: maintain two halves of the data
# max_heap: lower half (stores negated values for max behavior)
# min_heap: upper half
# Invariant: len(max_heap) == len(min_heap) or len(max_heap) == len(min_heap) + 1
# Invariant: max(max_heap) <= min(min_heap)
# Median:
# odd count: max_heap[0] (top of lower half)
# even count: average of tops of both halves
print('Two-heap strategy for O(log n) insert, O(1) median')Two-Heap MedianFinder Implementation
Maintain a max-heap for the lower half and a min-heap for the upper half. Always ensure the max-heap has the same size or one more element than the min-heap. When adding a number: push to max-heap, then balance by moving the max-heap's top to the min-heap if the top exceeds the min-heap's minimum, and rebalance sizes if needed.
import heapq
class MedianFinder:
def __init__(self):
self.lo = [] # max-heap (negated) for lower half
self.hi = [] # min-heap for upper half
def addNum(self, num):
heapq.heappush(self.lo, -num) # push to lower half
# Ensure max of lower <= min of upper
if self.hi and -self.lo[0] > self.hi[0]:
heapq.heappush(self.hi, -heapq.heappop(self.lo))
# Balance sizes: lo can have at most 1 more than hi
if len(self.lo) > len(self.hi) + 1:
heapq.heappush(self.hi, -heapq.heappop(self.lo))
elif len(self.hi) > len(self.lo):
heapq.heappush(self.lo, -heapq.heappop(self.hi))
def findMedian(self):
if len(self.lo) > len(self.hi):
return -self.lo[0] # odd count: top of lower half
return (-self.lo[0] + self.hi[0]) / 2
mf = MedianFinder()
for n in [1, 2, 3, 4, 5]: mf.addNum(n)
print(mf.findMedian()) # 3.0Trace Through MedianFinder Steps
Understanding why the two-heap invariant is maintained is critical for explaining the solution in an interview. Let's trace adding [5, 15, 1, 3] step by step. After each insertion: balance so the lower max-heap holds the smaller half. The invariant ensures max(lo) <= min(hi) always holds, which makes the median trivially accessible at the top of one or both heaps.
import heapq
# Manual trace for [5, 15, 1, 3]:
# add 5: lo=[-5] hi=[] median=5
# add 15: lo=[-5] hi=[15] median=(5+15)/2=10
# add 1: lo=[-5,-1] hi=[15] median=5
# add 3: lo=[-5,-3,-1] hi=[15] -- lo too big
# -> lo=[-5,-3] hi=[1,15] -- wait, wrong direction
# Actually:
# add 1: push to lo -> lo=[-5,-1], then 1>lo? No, -lo[0]=5>15? No
# lo has 2, hi has 1: balance -> move lo top to hi
# lo=[-1], hi=[5,15]
# Median = (-lo[0] + hi[0])/2 = (1+5)/2 = 3
mf2 = MedianFinder()
for n, expected in [(5, 5.0), (15, 10.0), (1, 5.0), (3, 4.0)]:
mf2.addNum(n)
print(f'After adding {n}: median={mf2.findMedian()} (expected ~{expected})')Sliding Window Median
The Sliding Window Median (LeetCode #480) is a harder variant: find the median of every window of size k as it slides across the array. The two-heap approach extends with a lazy deletion set to handle elements sliding out of the window. When an element leaves the window, mark it in the deletion set; when it reaches the top of either heap, discard it.
import heapq
def median_sliding_window(nums, k):
lo = [] # max-heap (negated)
hi = [] # min-heap
removed = {}
result = []
def balance():
# Move valid tops to correct side
while lo and removed.get(-lo[0], 0) > 0:
removed[-lo[0]] -= 1; heapq.heappop(lo)
while hi and removed.get(hi[0], 0) > 0:
removed[hi[0]] -= 1; heapq.heappop(hi)
for i, num in enumerate(nums):
heapq.heappush(lo, -num)
heapq.heappush(hi, -heapq.heappop(lo))
if len(hi) > len(lo): heapq.heappush(lo, -heapq.heappop(hi))
if i >= k:
out = nums[i - k]
removed[out] = removed.get(out, 0) + 1
balance()
if len(lo) > len(hi): heapq.heappush(hi, -heapq.heappop(lo))
if i >= k - 1:
if len(lo) > len(hi): result.append(float(-lo[0]))
else: result.append((-lo[0] + hi[0]) / 2.0)
return result
print(median_sliding_window([1,3,-1,-3,5,3,6,7], 3)) # [1,-1,-1,3,5,6]K-Way Merge: The Problem
Merge K Sorted Lists (LeetCode #23) is a fundamental problem with applications in external sorting, database merges, and distributed systems. Given k sorted linked lists totalling n nodes, merge them into one sorted list. The naive approach (merge two at a time) is O(kn) or O(n log k) with divide-and-conquer. The heap approach processes each node exactly once with O(log k) work per node: O(n log k) total.
class ListNode:
def __init__(self, val=0, next=None):
self.val = val
self.next = next
# Build a linked list from a Python list
def build_list(arr):
dummy = ListNode(0)
curr = dummy
for val in arr:
curr.next = ListNode(val)
curr = curr.next
return dummy.next
# Convert linked list to Python list for printing
def to_list(head):
result = []
while head:
result.append(head.val)
head = head.next
return result
print('K-way merge: O(n log k) using a min-heap of k heads')K-Way Merge with a Min-Heap
Initialize the heap with the first node of each list. At each step, pop the minimum, add it to the result, and push the next node from that list (if any). The heap always has at most k elements — one head per active list. Since we process n nodes total with O(log k) heap operations each, total time is O(n log k) and space is O(k) for the heap.
import heapq
def merge_k_lists(lists):
dummy = ListNode(0)
curr = dummy
heap = []
for i, node in enumerate(lists):
if node:
heapq.heappush(heap, (node.val, i, node))
while heap:
val, i, node = heapq.heappop(heap)
curr.next = node
curr = curr.next
if node.next:
heapq.heappush(heap, (node.next.val, i, node.next))
return dummy.next
lists = [
build_list([1, 4, 5]),
build_list([1, 3, 4]),
build_list([2, 6])
]
result = merge_k_lists(lists)
print(to_list(result)) # [1, 1, 2, 3, 4, 4, 5, 6]Smallest Range Covering K Lists
Smallest Range (LeetCode #632) finds the smallest range [lo, hi] such that at least one element from each of the k sorted lists lies within the range. Use a min-heap initialised with the first element of each list and track the current maximum. Shrink the range by always advancing the list with the current minimum. Stop when any list is exhausted.
import heapq
def smallest_range(nums):
heap = []
current_max = float('-inf')
for i, lst in enumerate(nums):
heapq.heappush(heap, (lst[0], i, 0))
current_max = max(current_max, lst[0])
best = [float('-inf'), float('inf')]
while heap:
current_min, list_idx, elem_idx = heapq.heappop(heap)
if current_max - current_min < best[1] - best[0]:
best = [current_min, current_max]
if elem_idx + 1 >= len(nums[list_idx]):
break # one list exhausted
next_val = nums[list_idx][elem_idx + 1]
heapq.heappush(heap, (next_val, list_idx, elem_idx + 1))
current_max = max(current_max, next_val)
return best
print(smallest_range([[4,10,15,24,26],[0,9,12,20],[5,18,22,30]]))
# [20, 24]K-th Smallest in a Matrix
Kth Smallest Element in a Sorted Matrix (LeetCode #378): an n×n matrix where each row and column is sorted. Find the kth smallest element. Treat each row as a sorted list and use k-way merge with a heap. Alternatively, binary search on the value range. The heap approach is O(k log n) which is efficient when k is small; binary search is O(n log(max-min)) which handles large k better.
import heapq
def kth_smallest_matrix(matrix, k):
n = len(matrix)
heap = [(matrix[0][0], 0, 0)]
count = 0
visited = {(0, 0)}
while heap:
val, r, c = heapq.heappop(heap)
count += 1
if count == k:
return val
# Push right neighbor
if c + 1 < n and (r, c+1) not in visited:
heapq.heappush(heap, (matrix[r][c+1], r, c+1))
visited.add((r, c+1))
# Push bottom neighbor
if r + 1 < n and (r+1, c) not in visited:
heapq.heappush(heap, (matrix[r+1][c], r+1, c))
visited.add((r+1, c))
return -1
matrix = [[1,5,9],[10,11,13],[12,13,15]]
print(kth_smallest_matrix(matrix, 8)) # 13Two Heaps for Running Statistics
The two-heap pattern generalises beyond median. You can use it to maintain a running quantile (e.g., the 25th percentile): size the lower heap to hold p*n elements and the upper heap to hold (1-p)*n elements. Each time an element is added, rebalance as before. This pattern appears in streaming statistics problems where you need efficient insertion and quantile queries simultaneously.
import heapq
# Generalised two-heap for arbitrary quantile p
# lo contains floor(p * count) elements
# hi contains the remaining elements
class QuantileFinder:
def __init__(self, p):
self.p = p # quantile (e.g., 0.5 for median)
self.lo = [] # max-heap
self.hi = [] # min-heap
self.count = 0
def add(self, num):
self.count += 1
heapq.heappush(self.lo, -num)
heapq.heappush(self.hi, -heapq.heappop(self.lo))
# Target: lo should have floor(p * count) elements
target_lo = int(self.p * self.count)
while len(self.lo) < target_lo:
heapq.heappush(self.lo, -heapq.heappop(self.hi))
while len(self.lo) > target_lo:
heapq.heappush(self.hi, -heapq.heappop(self.lo))
def quantile(self):
return -self.lo[0] if self.lo else self.hi[0]
qf = QuantileFinder(0.5) # median
for n in [1, 2, 3, 4, 5, 6]: qf.add(n)
print(qf.quantile()) # 3 (median of 1-6)Find K Closest Points to Origin
K Closest Points to Origin (LeetCode #973) uses a max-heap of size k. Push each point's squared distance (to avoid sqrt). When the heap exceeds k, pop the farthest. The remaining k points are the k closest. This is O(n log k). An alternative uses quickselect for O(n) average, but the heap solution is simpler to implement correctly and explain during an interview.
import heapq
def k_closest(points, k):
heap = [] # max-heap via negation
for x, y in points:
dist_sq = x*x + y*y
heapq.heappush(heap, (-dist_sq, x, y))
if len(heap) > k:
heapq.heappop(heap) # remove farthest
return [[x, y] for _, x, y in heap]
points = [[1,3], [-2,2], [5,8], [0,1], [-1,-1]]
print(k_closest(points, 2))
# Two closest to origin: [0,1] (dist=1) and [-1,-1] (dist=2)
# Verify by distances:
for x, y in points:
print(f'({x},{y}): dist^2 = {x*x+y*y}')Two Heaps: Time and Space Analysis
The two-heap approach for median achieves O(log n) per addNum and O(1) findMedian. Space is O(n) to store all elements. The k-way merge is O(n log k) time and O(k) space for the heap. These are near-optimal: you can prove a comparison-based lower bound of Omega(n log k) for k-way merge, showing the heap solution is asymptotically optimal. Always state these complexities clearly in interviews.
# Complexity summary for heap applications:
# Problem | Time per op | Space
# ----------------------|--------------|------
# MedianFinder.addNum | O(log n) | O(n)
# MedianFinder.find | O(1) | -
# Merge k sorted lists | O(n log k) | O(k)
# Kth smallest matrix | O(k log n) | O(n)
# K closest points | O(n log k) | O(k)
# Task scheduler | O(n log 26) | O(26)
# Kth largest stream | O(log k) | O(k)
# Sliding window median | O(n log k) | O(k)
print('Heap problems: identify k (heap size) vs n (input size)')Quick Check
Test your understanding of Data Structures & Algorithms — Coding Interview Prep concepts from this lesson.
Lesson Recap
In this lesson you learned: two-heap MedianFinder achieving O(log n) insert and O(1) median, k-way merge with a min-heap in O(n log k) time and O(k) space, and extensions including sliding window median, smallest range, and k-closest points. Next up we explore graph representations and traversal setup.
Frequently asked questions
Is the “Median from Data Stream and K-Way Merge” lesson free?
Yes — the full text of “Median from Data Stream and K-Way Merge” is free to read here on the web, and the DSA Interview Prep course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the DSA Interview Prep course, upgrade to CoddyKit PRO.
What will I learn in “Median from Data Stream and K-Way Merge”?
Maintain two heaps (max-heap of small half, min-heap of large half) for O(log n) median updates, and merge k sorted lists using a heap. You practise DSA Interview Prep with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.
Do I need any experience to start DSA Interview Prep?
No prior experience is required. DSA Interview Prep on CoddyKit is structured for beginners through advanced learners; this is — lesson 4 of 4, so you can start here or from the beginning and move at your own pace.
How long does the “Median from Data Stream and K-Way Merge” lesson take?
Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.
Can I write and run code in this DSA Interview Prep lesson?
Yes. Every DSA Interview Prep lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.
All lessons in this course
- Heap Property and Array Representation
- Heapify, Push, and Pop from Scratch
- Python heapq and Max-Heap Tricks
- Median from Data Stream and K-Way Merge