寻找长度为K的最小中位数子数组——求O(n log k)复杂度解法

问题描述 #

给定长度为N的整数数组，寻找长度为K（K为奇数）的子数组中中位数最小的那个。

示例：数组[1,3,3,2,1]，K=3，所有窗口分别为[1,3,3]、[3,3,2]、[3,2,1]，排序后中位数依次为3、3、2，答案是2。

现有解法问题分析 #

1. 滑动窗口+每次排序：结果正确但效率极低，时间复杂度为O(nk log k)，仅能拿到部分分数。
2. bisect模块实现：用insort维护有序窗口，但del操作的时间复杂度为O(k)，导致整体复杂度退化为O(nk)，无法满足高效需求。

你给出的bisect实现代码：

from bisect import insort, bisect_left
from typing import Sequence

def min_median(s: Sequence[int], m: int) -> int:
    n = len(s)
    result = 9**30
    window = sorted(s[:m])
    mid = m // 2

    result = window[mid]
    
    for i in range(m, n):
        insort(window, s[i])
        del window[bisect_left(window, s[i - m])]
        
        result = min(result, window[mid])
    
    return result

O(n log k) 解决方案 #

要达到目标复杂度，核心是用两个堆（大顶堆+小顶堆）+ 延迟删除的方案维护窗口元素，避免直接操作有序数组的高开销：

• 大顶堆（用负号模拟）：存储窗口中较小的一半元素，堆顶为这部分的最大值，也就是当前窗口的中位数
• 小顶堆：存储窗口中较大的一半元素，堆顶为这部分的最小值
• 延迟删除哈希表：记录待删除元素的计数，避免直接删除堆元素的O(k)操作，后续堆顶元素失效时再清理

实现逻辑 #

1. 初始化窗口：将前K个元素分配到两个堆，保证大顶堆的元素数量比小顶堆多1（因为K是奇数），此时大顶堆堆顶就是初始窗口的中位数。
2. 滑动窗口维护：
   • 标记窗口左端元素为待删除
   • 将新元素加入对应堆
   • 清理两个堆的失效元素（已标记待删除的堆顶）
   • 调整堆的大小平衡，确保大顶堆始终比小顶堆多1个元素
3. 更新最小中位数：每次窗口平衡后，取大顶堆堆顶作为当前窗口中位数，更新全局最小值。

代码实现 #

import heapq
from collections import defaultdict

def min_median(nums, k):
    max_heap = []  # 用负号模拟大顶堆，存储较小的一半元素
    min_heap = []  # 小顶堆，存储较大的一半元素
    delayed = defaultdict(int)  # 延迟删除的计数表
    n = len(nums)
    min_result = float('inf')
    required_max_size = (k + 1) // 2  # 大顶堆需要维持的元素数量

    # 初始化前k个元素到堆中
    for num in nums[:k]:
        heapq.heappush(max_heap, -num)
    # 调整堆大小，让大顶堆保留required_max_size个元素
    for _ in range(k - required_max_size):
        val = -heapq.heappop(max_heap)
        heapq.heappush(min_heap, val)
    
    min_result = -max_heap[0]

    # 滑动窗口处理后续元素
    for i in range(k, n):
        left_val = nums[i - k]
        # 标记左端元素为待删除
        if left_val <= -max_heap[0]:
            delayed[left_val] += 1
        else:
            delayed[left_val] += 1
        
        # 添加当前元素到对应堆
        current_val = nums[i]
        if current_val <= -max_heap[0]:
            heapq.heappush(max_heap, -current_val)
        else:
            heapq.heappush(min_heap, current_val)
        
        # 清理大顶堆的失效元素
        while max_heap and delayed[-max_heap[0]] > 0:
            delayed[-max_heap[0]] -= 1
            heapq.heappop(max_heap)
        
        # 清理小顶堆的失效元素
        while min_heap and delayed[min_heap[0]] > 0:
            delayed[min_heap[0]] -= 1
            heapq.heappop(min_heap)
        
        # 调整堆的大小平衡
        while len(max_heap) > required_max_size:
            val = -heapq.heappop(max_heap)
            heapq.heappush(min_heap, val)
            # 清理可能的失效元素
            while max_heap and delayed[-max_heap[0]] > 0:
                delayed[-max_heap[0]] -= 1
                heapq.heappop(max_heap)
        while len(max_heap) < required_max_size:
            val = heapq.heappop(min_heap)
            heapq.heappush(max_heap, -val)
            # 清理可能的失效元素
            while min_heap and delayed[min_heap[0]] > 0:
                delayed[min_heap[0]] -= 1
                heapq.heappop(min_heap)
        
        # 更新最小中位数
        min_result = min(min_result, -max_heap[0])
    
    return min_result

# 测试示例
print(min_median([1,3,3,2,1], 3))  # 输出2

复杂度说明 #

• 每个元素最多经历两次堆操作（入堆和出堆），每次堆操作的时间复杂度为O(log k)，整体时间复杂度为O(n log k)
• 空间复杂度为O(k)，用于存储两个堆和延迟删除哈希表

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