数据流的中位数 | 秋招知识库

比如现在有 $6$ 个数： $1, 5, 6, 2, 3, 4$ ，要计算中位数，可以把这 $6$ 个数从小到大排序，得到 $1, 2, 3, 4, 5, 6$ ，中间 $3$ 和 $4$ 的平均值 $3.5$ 就是中位数。

回顾一下百科中关于中位数的定义：

中位数……可将数值集合划分为相等的两部分。

中位数把这 $6$ 个数均分成了左右两部分，小的那一组记作 $left = [1, 2, 3]$ ，大的那一组记作 $right = [4, 5, 6]$ 。我们要计算的中位数，就来自 $left$ 中的最大值，以及 $right$ 中的最小值。

随着 $addNum$ 不断地添加数字，我们需要：

保证 $left$ 的大小和 $right$ 的大小尽量相等。规定：在有奇数个数时， $left$ 比 $right$ 多 $1$ 个数。
保证 $left$ 的所有元素都小于等于 $right$ 的所有元素。

只要时时刻刻满足以上两个要求（满足中位数的定义），我们就可以用 $left$ 中的最大值以及 $right$ 中的最小值计算中位数。

分类讨论：

如果当前 $left$ 的大小和 $right$ 的大小相等：
- 如果添加的数字 $num$ 比较大，比如添加 $7$ ，那么把 $7$ 加到 $right$ 中。现在 $left$ 比 $right$ 少 $1$ 个数，不符合前文的规定，所以必须把 $right$ 的最小值从 $right$ 中去掉，添加到 $left$ 中。如此操作后，可以保证 $left$ 的所有元素都小于等于 $right$ 的所有元素。
- 如果添加的数字 $num$ 比较小，比如添加 $0$ ，那么把 $0$ 加到 $left$ 中。
- 这两种情况可以合并：无论 $num$ 是大是小，都可以先把 $num$ 加到 $right$ 中，然后把 $right$ 的最小值从 $right$ 中去掉，并添加到 $left$ 中。
如果当前 $left$ 比 $right$ 多 $1$ 个数：
- 如果添加的数字 $num$ 比较大，比如添加 $7$ ，那么把 $7$ 加到 $right$ 中。
- 如果添加的数字 $num$ 比较小，比如添加 $0$ ，那么把 $0$ 加到 $left$ 中。现在 $left$ 比 $right$ 多 $2$ 个数，不符合前文的规定，所以必须把 $left$ 的最大值从 $left$ 中去掉，添加到 $right$ 中。如此操作后，可以保证 $left$ 的所有元素都小于等于 $right$ 的所有元素。
- 这两种情况可以合并：无论 $num$ 是大是小，都可以先把 $num$ 加到 $left$ 中，然后把 $left$ 的最大值从 $left$ 中去掉，并添加到 $right$ 中。

最后，我们需要什么样的数据结构？这个数据结构要能高效地执行如下操作：

添加元素。
找到最大（小）值。
删除最大（小）值。

这个数据结构是堆。

$left$ 是最大堆， $right$ 是最小堆。

如果当前有奇数个元素，中位数是 $left$ 的堆顶。
如果当前有偶数个元素，中位数是 $left$ 的堆顶和 $right$ 的堆顶的平均值。

Python 解法1C++ 解法1

python

class MedianFinder:
    def __init__(self):
        self.left = []  # 最大堆
        self.right = []  # 最小堆

    def addNum(self, num: int) -> None:
        if len(self.left) == len(self.right):
            heappush_max(self.left, heappushpop(self.right, num))
        else:
            heappush(self.right, heappushpop_max(self.left, num))

    def findMedian(self) -> float:
        if len(self.left) > len(self.right):
            return self.left[0]
        return (self.left[0] + self.right[0]) / 2

cpp

// C++ 版待补充

C++ 解法2

cpp

class MedianFinder {
    priority_queue<int> left; // 最大堆
    priority_queue<int, vector<int>, greater<>> right; // 最小堆

public:
    void addNum(int num) {
        if (left.size() == right.size()) {
            right.push(num);
            left.push(right.top());
            right.pop();
        } else {
            left.push(num);
            right.push(left.top());
            left.pop();
        }
    }

    double findMedian() {
        if (left.size() > right.size()) {
            return left.top();
        }
        return (left.top() + right.top()) / 2.0;
    }
};