0Pricing
Coding Interview Prep · 课时

数据流中位数与 K 路合并

维护两个堆(较小一半使用大根堆,较大一半使用小根堆),以 O(log n) 的时间更新中位数,并使用堆合并 k 个有序列表。

数据流中位数与 K 路合并 是 CoddyKit 上的免费 Coding Interview Prep 课时。 这是第 4 节课,共 4 节。 你可以在下方免费阅读本课时的完整内容 — 然后在浏览器中使用内置代码编辑器和全天候 AI 导师进行实践。 这是 Coding Interview Prep 学习路径的一部分,你的进度在网页和 CoddyKit 应用中同步。 Coding Interview Prep 课程共包含 4 节课。

数据流中的中位数问题

查找数据流中的中位数(LeetCode #295)要求您高效支持两种操作:使用 addNum(num) 添加数字,以及使用 findMedian() 返回当前中位数。偶数长度列表的中位数是中间两个值的平均值。使用暴力方法维护有序列表时,插入复杂度为 O(n),获取中位数的复杂度为 O(1)。最优解使用两个堆,将插入复杂度降为 O(log n),同时获取中位数仍为 O(1)。

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')

使用两个堆实现 MedianFinder

维护一个存放较小一半元素的最大堆和一个存放较大一半元素的最小堆。始终确保最大堆的大小与最小堆相同,或只比最小堆多一个元素。添加数字时:先将其推入最大堆;如果最大堆的堆顶大于最小堆的最小值,就将最大堆的堆顶移入最小堆,以此进行平衡;如有需要,再平衡两个堆的大小。

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.0

跟踪 MedianFinder 的执行步骤

理解为何能够维持两个堆的不变量,对于在面试中解释该解法至关重要。下面逐步跟踪添加 [5, 15, 1, 3] 的过程。每次插入后都进行平衡,使存放较小元素一半的最大堆保持正确状态。该不变量确保 max(lo) <= min(hi) 始终成立,因此可以直接从一个堆或两个堆的堆顶获取中位数。

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})')

滑动窗口中位数

滑动窗口中位数(LeetCode #480)是一个更复杂的变体:当大小为 k 的窗口在数组上滑动时,找出每个窗口的中位数。两个堆的方法可以结合延迟删除集合,以处理滑出窗口的元素。当元素离开窗口时,将其标记在删除集合中;当它到达任一堆的堆顶时,再将其丢弃。

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 路归并:问题

合并 k 个有序列表(LeetCode #23)是一个基础问题,应用于外部排序、数据库合并和分布式系统。给定 k 个有序链表,总共有 n 个节点,请将它们合并为一个有序列表。朴素方法是每次合并两个列表,复杂度为 O(kn);使用分治法时,复杂度为 O(n log k)。堆方法对每个节点只处理一次,每个节点执行 O(log k) 的工作,总复杂度为 O(n log k)。

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 路归并

将每个列表的第一个节点初始化到堆中。每一步都弹出最小节点,将其加入结果中,然后将该列表的下一个节点(如果存在)推入堆中。堆中始终至多有 k 个元素,即每个活跃列表的一个头节点。总共处理 n 个节点,每个节点执行 O(log k) 的堆操作,因此总时间复杂度为 O(n log k),堆的空间复杂度为 O(k)。

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]

覆盖 k 个列表的最小范围

最小范围(LeetCode #632)要找出最小范围 [lo, hi],使得 k 个有序列表中的每个列表至少有一个元素位于该范围内。使用每个列表的第一个元素初始化最小堆,并跟踪当前最大值。每次推进当前最小值所在的列表,以不断缩小范围。当任意列表耗尽时停止。

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 小元素

有序矩阵中的第 k 小元素(LeetCode #378):给定一个 n×n 矩阵,其中每一行和每一列都已排序,找出第 k 小的元素。将每一行视为有序列表,并使用堆进行 k 路归并。另一种方法是在取值范围上进行二分查找。堆方法的复杂度为 O(k log n),当 k 较小时效率较高;二分查找的复杂度为 O(n log(max-min)),处理较大的 k 时表现更好。

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))  # 13

使用两个堆维护动态统计量

两个堆的模式不只适用于中位数。您可以使用它维护动态分位数(例如第 25 百分位数):让较小元素堆的大小为 p*n,让较大元素堆的大小为 (1-p)*n。每次添加元素时,像之前一样重新平衡。这种模式出现在流式统计问题中,适用于同时需要高效插入和分位数查询的场景。

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)

查找距离原点最近的 k 个点

距离原点最近的 k 个点(LeetCode #973)使用大小为 k 的最大堆。推入每个点的距离平方(以避免计算平方根)。当堆的大小超过 k 时,弹出距离最远的点。剩下的 k 个点就是距离最近的 k 个点。该方法的复杂度为 O(n log k)。另一种方法是使用快速选择,平均复杂度为 O(n),但堆方法更容易正确实现,也更容易在面试中解释。

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}')

两个堆:时间与空间分析

使用两个堆处理中位数时,每次 addNum 的复杂度为 O(log n),findMedian 的复杂度为 O(1)。存储所有元素需要 O(n) 的空间。k 路归并的时间复杂度为 O(n log k),堆的空间复杂度为 O(k)。这些复杂度接近最优:可以证明,基于比较的 k 路归并下界为 Omega(n log k),这说明堆解法在渐进意义下是最优的。在面试中,请始终清楚地说明这些复杂度。

# 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)')

快速检查

测试您对本课中数据结构与算法——编程面试准备相关概念的理解。

课程回顾

在本课中,您学习了:达到 O(log n) 插入和 O(1) 获取中位数的两个堆 MedianFinder,时间复杂度为 O(n log k)、空间复杂度为 O(k) 的最小堆k 路归并,以及包括滑动窗口中位数、最小范围和距离最近的 k 个点在内的扩展内容。接下来,我们将探索图的表示方式和遍历准备工作。

常见问题解答

「数据流中位数与 K 路合并」课时是免费的吗?

是的 — 「数据流中位数与 K 路合并」的完整文本可在网页上免费阅读。要进行交互式练习(内置代码编辑器和全天候 AI 导师)并解锁 Coding Interview Prep 课程的其余内容,请升级到 CoddyKit PRO。 Coding Interview Prep 课程共包含 4 节课。

「数据流中位数与 K 路合并」这节课中我会学到什么?

维护两个堆(较小一半使用大根堆,较大一半使用小根堆),以 O(log n) 的时间更新中位数,并使用堆合并 k 个有序列表。 你通过在浏览器中直接运行的动手代码来练习 Coding Interview Prep,全天候 AI 导师会在你学习这节课的过程中回答你的问题。

学习 Coding Interview Prep 需要有经验吗?

无需任何先前经验。CoddyKit 上的 Coding Interview Prep 课程适合初学者到高级学习者,你可以从这里开始或从头开始,按照自己的节奏学习。 这是第 4 节课,共 4 节。

「数据流中位数与 K 路合并」课时需要多长时间?

大多数 CoddyKit 课程大约需要 5–10 分钟。每节课都很精短且互动,所以你能稳步进步,并在网页和应用中从离开的地方继续。

我能在这节 Coding Interview Prep 课中编写并运行代码吗?

能。每节 Coding Interview Prep 课都包含内置代码编辑器,你可以在浏览器中直接编写并运行真实代码,并获得即时 AI 反馈 — 无需本地设置。

此课程中的所有课时

  1. 堆性质与数组表示
  2. 从零实现建堆、推入与弹出
  3. Python heapq 与大根堆技巧
  4. 数据流中位数与 K 路合并
← 返回 Coding Interview Prep