題目網址：https://leetcode.cn/problems/find-median-from-data-stream/

題意：中位數是指有序 array 裡中間的數。若 array 長度偶數, 則中位數是中間兩數的平均值。

  • arr = [2,3,4], 則中位數 = 3
  • arr = [2,3], 則中位數 = (2+3) / 2 = 2.5

請實現 MedianFinder class：

  • MedianFinder()：初始化 MedianFinder
  • void addNum(int num)：新增 num 到資料結構中
  • double findMedian()：找出目前的中位數

Solution 1：

想法：利用 Balanced BST(紅黑樹), C++ 可直接使用 multiset, 將元素 push 到 multiset 中它會自動排好序。我們用一個 iterator mid 指向目前的中位數（若 multiset 中的元素個數為偶數, 則讓 mid 指向中間靠左的元素）

  • 若 num 加入前, set 中的元素為偶數個(OOXOOO, 其中 X 代表 mid), 則加入 nums 後, 分兩種情況:
    • 若 nums >= *mid : OOXOOO(加入前) ➔ OOXOOOO(加入後), 則要將 mid 往右移一位, OOXOOOO ➔ OOOXOOO
    • 若 nums < *mid : OOXOOO(加入前) ➔ OOOXOOO(加入後), 則 mid 不需移動
    • 加入 num 後, set 中的元素變為奇數個, 因此返回 mid 即可
  • 若 num 加入前, set 中的元素為奇數個(OOXOO), 則加入 nums 後, 分兩種情況:
    • 若 nums >= *mid : OOXOO(加入前) ➔ OOXOOO(加入後), 則 mid 不需移動
    • 若 nums < *mid : OOXOO(加入前) ➔ OOOXOO(加入後), 則要將 mid 往左移一位, OOOXOO ➔ OOXOOO
    • 加入 num 後, set 中的元素變為偶數個, 因此要返回中間兩數的平均值
class MedianFinder {
public:
    MedianFinder() {}

    void addNum(int num) {
        set.emplace(num);

        // 若為第一個元素, 則初始化 mid
        if (set.size() == 1) {
            mid = set.begin();
            return;
        }

        // 加入 num 前為偶數個, OOXOOO
        if (set.size() % 2 == 1) {
            if (num >= *mid) {
                mid = next(mid, 1);
            }
        } else { // 加入 num 前為奇數個, OOXOO
            if (num < *mid) {
                mid = prev(mid, 1);
            }
        }
    }

    double findMedian() {
         // 若 set.size() 為偶數, 則中位數 = (左中位數 + 右中位數) / 2, 要避免 overflow
        return *mid * 0.5 + *next(mid, (set.size() + 1) % 2) * 0.5;
    }

private:
    multiset<int> set;
    multiset<int>::iterator mid;
};
  • time：
    • addNum()：$O(log(n))$, Balanced BST 新增元素所需的時間複雜度為 $O(h)$, 因為是 Balanced, 所以 $h = log(n)$
    • findMedian()：$O(1)$
  • space：$O(n)$ ➔ set

Solution 2：

想法：利用 Heap, 使用兩個 heap 來分別記錄中間左中位數、右中位數, 因為當 set 中元素為偶數時, 我們需要取得這兩個數來計算平均值

  • Max Heap smaller 紀錄所有 ≤ 中位數的數, smaller.top() 可得到所有 ≤ 中位數的數 中最大的值, 也就是左中位數
  • Min Heap larger 紀錄所有 > 中位數的數, larger.top() 可得到所有 > 中位數的數 中最小的值, 也就是右中位數

在新增 num 時, 我們必須使兩個 heap 的長度保持平衡, 平衡的定義如下:

  • 滿足 0 ≤ | smaller.size() - larger.size() | ≤ 1

當兩個 heap 的長度不平衡時, 則把長度較長者的 top 給 pop 掉, 並把其 push 到另一個 heap

  • 要返回中位數時, 我們比較兩個 heap 的長度

    • 若 smaller.size() > larger.size() 時, 返回 smaller.top()

      e.g. smaller = [1,2,3], larger = [4,5], 則中位數 = 3

    • 若 larger.size() > smaller.size() 時, 返回 larger.top()

      e.g. smaller = [1,2], larger = [3,4,5], 則中位數 = 3

    • 若 smaller.size() == larger.size(), 代表 set 中元素個數為偶數, 返回中間兩數的平均值, 也就是 (smaller.top() + larger.top()) / 2.0

      e.g. smaller = [1,2,3], larger = [4,5,6], 則中位數 = (3 + 4) / 2 = 3.5

class MedianFinder {
public:
    MedianFinder() {}

    void addNum(int num) {
        // 將 num 加入到 heap 中
        if (smaller.empty() || num <= smaller.top()) {
            smaller.emplace(num);
        } else {
            larger.emplace(num);
        }

        const int m = smaller.size(), n = larger.size();

        // 平衡 smaller 和 larger 的 size
        if (m - n > 1) {
            larger.emplace(smaller.top());
            smaller.pop();
        } else if (n - m > 1) {
            smaller.emplace(larger.top());
            larger.pop();
        }
    }

    double findMedian() {
        if (smaller.size() > larger.size()) { // s = [1,2,3], l = [4,5]
            return smaller.top();
        }

        if (larger.size() > smaller.size()) { // s = [1,2], l = [3,4,5]
            return larger.top();
        }

        return smaller.top() * 0.5 + larger.top() * 0.5; // 避免 overflow 的寫法
    }

private:
    priority_queue<int> smaller; // max heap
    priority_queue<int, vector<int>, greater<int>> larger; // min heap
};
  • time：
    • addNum()：$O(log(n))$ 從 heap 中新增 or 刪除元素
    • findMedian()：$O(1)$
  • space：$O(n)$ ➔ smaller + larger 的長度不超過 n