数据流的中位数

题意

设计数据结构，支持动态添加数字和查询中位数。

思路

用两个堆维护数据流：

lo（大顶堆）：存较小的一半，堆顶是这半的最大值。
hi（小顶堆）：存较大的一半，堆顶是这半的最小值。

插入时先放入 lo，再弹出 lo 堆顶放入 hi（保证 hi 中所有值 $\geqslant$ lo 中所有值）。若 lo 比 hi 少，则从 hi 弹回 lo（保证 lo 元素数 $\geqslant$ hi）。

两个不变式：lo 中所有值 $\leqslant$ hi 中所有值；lo 元素数比 hi 多 0 或 1。

查询中位数：lo 多时取 lo 堆顶，否则取两堆顶平均值。

代码

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42#include <bits/stdc++.h>
using namespace std;

class MedianFinder {
    // lo 保存较小的一半，且元素个数始终等于 hi 或比 hi 多 1。
    priority_queue<int> lo;
    priority_queue<int, vector<int>, greater<>> hi;

public:
    void addNum(int num) {
        lo.push(num);
        hi.push(lo.top());
        lo.pop();
        if (lo.size() < hi.size()) {
            lo.push(hi.top());
            hi.pop();
        }
    }

    double findMedian() {
        return lo.size() > hi.size() ? lo.top() : (lo.top() + hi.top()) / 2.0;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int q;
    cin >> q;
    MedianFinder mf;
    while (q--) {
        string op;
        cin >> op;
        if (op == "add") {
            int v;
            cin >> v;
            mf.addNum(v);
        } else
            cout << fixed << setprecision(1) << mf.findMedian() << ' ';
    }
    return 0;
}

python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32#!/usr/bin/env python3
import heapq


class MedianFinder:
    def __init__(self):
        self.lo = []
        self.hi = []

    def addNum(self, n):
        heapq.heappush(self.lo, -n)
        heapq.heappush(self.hi, -heapq.heappop(self.lo))
        if len(self.lo) < len(self.hi):
            heapq.heappush(self.lo, -heapq.heappop(self.hi))

    def findMedian(self):
        return float(-self.lo[0]) if len(self.lo) > len(self.hi) else (-self.lo[0] + self.hi[0]) / 2


def main():
    q = int(input())
    mf = MedianFinder()
    for _ in range(q):
        op = input().split()
        if op[0] == "add":
            mf.addNum(int(op[1]))
        else:
            print(mf.findMedian(), end=" ")


if __name__ == "__main__":
    main()

复杂度

时间复杂度：addNum $O(\log n)$ ，findMedian $O(1)$ 。
空间复杂度： $O(n)$ 。

总结

双堆中位数的关键是两个不变式：大小关系和元素数关系。插入时的"过一遍对方堆"保证大小关系，"多则弹回"保证元素数关系。Python 中 lo 用负数模拟大顶堆。

目录

题意

思路

代码

复杂度

总结