Python中高效计算流式单数字序列中位数的优化方案求助

问题描述 #

给定一串单数字组成的字符串，数字按队列顺序逐个进入空房间：

• 每分钟从左侧取一个数字进入房间
• 每次新数字进入后，需记录当前房间内所有数字的中位数（升序排列后，奇数个取中间数；偶数个取两个中间数中较小的）
• 最终将所有记录的中位数按顺序拼接成字符串输出

示例：
输入：21423814127333
最终输出：21222222222233

原方案问题分析 #

你当前的代码每次添加元素后都调用sort()，时间复杂度为O(n²logn)。当输入字符串长度很大时（比如上万甚至更长的测试用例），反复排序的开销会导致超时。

优化方案：双堆法 #

使用两个堆来维护当前数字的前后两半，实现O(logn)时间复杂度的插入和中位数查询：

• 最大堆：存储较小的一半数字（用Python的heapq模拟，存负数实现最大堆功能），堆顶是这部分的最大值，也就是当前的中位数候选
• 最小堆：存储较大的一半数字，堆顶是这部分的最小值
• 维护规则：最大堆的大小要么等于最小堆，要么比最小堆大1。这样无论当前元素总数是奇数还是偶数，中位数都是最大堆的堆顶（偶数个时取较小的中位数，正好是最大堆的堆顶）

优化后代码 #

import heapq

num = input().strip()
result = []

# 最大堆（存负数）和最小堆
max_heap = []
min_heap = []

for digit in num:
    d = int(digit)
    # 先插入到对应的堆
    if not max_heap or d <= -max_heap[0]:
        heapq.heappush(max_heap, -d)
    else:
        heapq.heappush(min_heap, d)
    
    # 平衡两个堆的大小
    if len(max_heap) > len(min_heap) + 1:
        # 把最大堆的堆顶移到最小堆
        val = -heapq.heappop(max_heap)
        heapq.heappush(min_heap, val)
    elif len(min_heap) > len(max_heap):
        # 把最小堆的堆顶移到最大堆
        val = heapq.heappop(min_heap)
        heapq.heappush(max_heap, -val)
    
    # 当前中位数是最大堆的堆顶（负数取反）
    result.append(str(-max_heap[0]))

print(''.join(result))

代码说明 #

1. 堆的初始化：max_heap用负数存储，模拟最大堆；min_heap直接用默认的最小堆
2. 插入逻辑：判断当前数字应该进入最大堆还是最小堆，保证堆内元素的有序性
3. 堆平衡：调整两个堆的大小，确保最大堆的大小始终满足len(max_heap) == len(min_heap)或len(max_heap) == len(min_heap)+1
4. 中位数获取：每次插入平衡后，直接取max_heap的堆顶（取反后）作为当前中位数，加入结果列表

该方案的时间复杂度为O(nlogn)，每个元素的插入和堆调整都是O(logn)的操作，能够高效处理大规模输入，解决超时问题。

内容的提问来源于stack exchange，提问作者Blueberry