Python中如何实现O(log n)时间复杂度支持增删的动态中位数求解

内置结构说明 #

Python标准库没有提供类似C++ multiset的内置平衡二叉搜索树结构，无法直接用内置组件一步实现需求，可通过以下两种方案实现。 #

方案1：纯标准库实现：懒删除双堆法 #

不需要依赖第三方库，基于heapq模块配合延迟删除逻辑实现，所有操作时间复杂度均为O(log n)。

核心原理 #

仍然沿用双堆的基础设计：

• 大顶堆max_heap存储数值较小的下半部分元素（Python的heapq是小顶堆，存储元素的相反数模拟大顶堆）
• 小顶堆min_heap存储数值较大的上半部分元素
  额外引入两个哈希表记录待删除元素的计数，不需要立即从堆中删除元素，只在堆顶元素需要被访问时才清理已标记删除的元素，避免堆的O(n)查找开销。

实现代码 #

import heapq
class DynamicMedian:
    def __init__(self):
        self.max_heap = []  # 存下半部分，元素为相反数，模拟大顶堆
        self.min_heap = []  # 存上半部分
        self.delayed_max = {}  # max_heap待删除元素计数
        self.delayed_min = {}  # min_heap待删除元素计数
        self.size_max = 0  # max_heap有效元素数
        self.size_min = 0  # min_heap有效元素数
    
    def _clean_max_top(self):
        # 清理max_heap堆顶的已删除元素
        while self.max_heap:
            current = -self.max_heap[0]
            if current in self.delayed_max and self.delayed_max[current] > 0:
                self.delayed_max[current] -= 1
                if self.delayed_max[current] == 0:
                    del self.delayed_max[current]
                heapq.heappop(self.max_heap)
            else:
                break
    
    def _clean_min_top(self):
        # 清理min_heap堆顶的已删除元素
        while self.min_heap:
            current = self.min_heap[0]
            if current in self.delayed_min and self.delayed_min[current] > 0:
                self.delayed_min[current] -= 1
                if self.delayed_min[current] == 0:
                    del self.delayed_min[current]
                heapq.heappop(self.min_heap)
            else:
                break
    
    def _balance(self):
        # 调整两个堆的大小：保证size_max = size_min 或者 size_max = size_min + 1
        while self.size_max > self.size_min + 1:
            # max_heap元素太多，移到min_heap
            self._clean_max_top()
            val = -heapq.heappop(self.max_heap)
            self.size_max -= 1
            heapq.heappush(self.min_heap, val)
            self.size_min += 1
        while self.size_min > self.size_max:
            # min_heap元素太多，移到max_heap
            self._clean_min_top()
            val = heapq.heappop(self.min_heap)
            self.size_min -= 1
            heapq.heappush(self.max_heap, -val)
            self.size_max += 1
    
    def add(self, x):
        self._clean_max_top()
        if not self.max_heap or x <= -self.max_heap[0]:
            heapq.heappush(self.max_heap, -x)
            self.size_max += 1
        else:
            heapq.heappush(self.min_heap, x)
            self.size_min += 1
        self._balance()
    
    def remove(self, x):
        self._clean_max_top()
        if x <= -self.max_heap[0]:
            # 元素在max_heap中
            self.delayed_max[x] = self.delayed_max.get(x, 0) + 1
            self.size_max -= 1
        else:
            # 元素在min_heap中
            self.delayed_min[x] = self.delayed_min.get(x, 0) + 1
            self.size_min -= 1
        self._balance()
    
    def get_median(self):
        self._clean_max_top()
        total = self.size_max + self.size_min
        if total % 2 == 1:
            return -self.max_heap[0]
        else:
            self._clean_min_top()
            return (-self.max_heap[0] + self.min_heap[0]) / 2

方案2：第三方库SortedList实现（代码更简洁） #

如果允许使用第三方库，sortedcontainers库提供的SortedList是纯Python实现的有序列表，底层基于分段平衡BST设计，插入、删除、按索引访问的时间复杂度均为O(log n)，完全匹配需求。

实现步骤 #

1. 安装依赖：pip install sortedcontainers
2. 代码实现：

from sortedcontainers import SortedList
class DynamicMedian:
    def __init__(self):
        self.sl = SortedList()
    
    def add(self, x):
        self.sl.add(x)
    
    def remove(self, x):
        self.sl.remove(x)
    
    def get_median(self):
        n = len(self.sl)
        if n % 2 == 1:
            return self.sl[n // 2]
        else:
            return (self.sl[n//2 - 1] + self.sl[n//2]) / 2

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