PAT 1057 Stack with PeekMedian：三種解法從暴力到樹狀數組再到雙 Multiset

題目大意

實現一個支援三種操作的堆疊：

• Push key：壓入鍵值 (1 ≤ key ≤ 10^5)
• Pop：彈出頂端元素並輸出，空則輸出 Invalid
• PeekMedian：輸出當前堆疊中位數 (偶數個取第 N/2 小，奇數取第 (N+1)/2 小)，空則 Invalid

N ≤ 10^5 操作，需高效處理中位數查詢。

解法演進

解法一：暴力排序 — TLE

每次 PeekMedian 複製堆疊到陣列、排序、取中位數。
時間：Push/Pop O(1)，PeekMedian O(N log N)，最壞 O(N^2 log N) → TLE。

解法二：樹狀數組 (Fenwick Tree) — AC

核心思想

• 鍵值範圍 [1, 10^5]，可作為樹狀數組索引。
• 樹狀數組維護 freq[x] = 當前堆疊中鍵值 x 的出現次數。
• Push：update(key, +1)；Pop：update(key, -1)。
• PeekMedian：尋找第 k = (size+1)/2 小元素 → 樹狀數組上二分查找前綴和 ≥ k 的最小索引。

關鍵操作

const int MAXV = 100000;
int BIT[MAXV + 2];

int lowbit(int x) { return x & -x; }

void update(int idx, int delta) {
    for (; idx <= MAXV; idx += lowbit(idx)) BIT[idx] += delta;
}

int query(int idx) { // 前綴和 [1..idx]
    int sum = 0;
    for (; idx > 0; idx -= lowbit(idx)) sum += BIT[idx];
    return sum;
}

// 二分查找第 k 小 (k 從 1 開始)
int findKth(int k) {
    int l = 1, r = MAXV, ans = MAXV;
    while (l <= r) {
        int mid = (l + r) >> 1;
        if (query(mid) >= k) {
            ans = mid;
            r = mid - 1;
        } else l = mid + 1;
    }
    return ans;
}

複雜度

• Push/Pop：O(log MAXV)
• PeekMedian：O(log MAXV) (二分查找 log MAXV 步，每步 query O(log MAXV) → 可優化為單次 O(log MAXV) 下降式查找)
• 總計 O(N log MAXV)，極快 AC。

解法三：雙 Multiset (平衡 BST) — AC，更優雅

核心思想

• 維護兩個有序多重集合：
• lower (大根堆序，greater<int>)：存較小一半元素，lower.size() == (n+1)/2。
• upper (小根堆序)：存較大一半元素。
• 中位數永遠是 *lower.begin() (lower 最大值)。

平衡邏輯

multiset<int, greater<int>> lower; // 降序，begin() 為最大
multiset<int> upper;                 // 升序，begin() 為最小

void adjust() {
    // 保證 lower.size() == upper.size() 或 lower.size() == upper.size() + 1
    if (upper.size() > lower.size()) {
        lower.insert(*upper.begin());
        upper.erase(upper.begin());
    } else if (lower.size() > upper.size() + 1) {
        upper.insert(*lower.begin());
        lower.erase(lower.begin());
    }
    mid = *lower.begin(); // 更新中位數快取
}

void push(int key) {
    if (key > mid) upper.insert(key);
    else lower.insert(key);
    adjust();
}

void pop(int key) { // key 為堆疊頂端實際值
    if (key > mid) upper.erase(upper.find(key));
    else lower.erase(lower.find(key));
    if (stack.empty()) mid = 0;
    else adjust();
}

優勢

• 無需鍵值範圍限制，適用於任意可比較型別。
• 程式碼直觀，標準庫 multiset 紅黑樹實現，常數略大但 N=10^5 完全足夠。
• 易於擴展：支援任意順序統計量 (第 k 小) 只需調整平衡條件。

完整 C++ 參考程式碼 (雙 Multiset 版)

#include <bits/stdc++.h>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    int N;
    cin >> N;
    stack st;
    multiset> lower; // 小半部，降序
    multiset upper;               // 大半部，升序
    int mid = 0; // 當前中位數
    
    string op;
    for (int i = 0; i < N; ++i) {
        cin >> op;
        if (op[1] == 'o') { // Pop
            if (st.empty()) {
                cout << "Invalid\n";
            } else {
                int key = st.top(); st.pop();
                cout << key << '\n';
                if (key > mid) upper.erase(upper.find(key));
                else lower.erase(lower.find(key));
                if (st.empty()) mid = 0;
                else {
                    // adjust
                    if (upper.size() > lower.size()) {
                        lower.insert(*upper.begin());
                        upper.erase(upper.begin());
                    } else if (lower.size() > upper.size() + 1) {
                        upper.insert(*lower.begin());
                        lower.erase(lower.begin());
                    }
                    mid = *lower.begin();
                }
            }
        } else if (op[1] == 'e') { // PeekMedian
            if (st.empty()) cout << "Invalid\n";
            else cout << mid << '\n';
        } else { // Push
            int key; cin >> key;
            st.push(key);
            if (st.size() == 1) mid = key;
            else {
                if (key > mid) upper.insert(key);
                else lower.insert(key);
                // adjust
                if (upper.size() > lower.size()) {
                    lower.insert(*upper.begin());
                    upper.erase(upper.begin());
                } else if (lower.size() > upper.size() + 1) {
                    upper.insert(*lower.begin());
                    lower.erase(lower.begin());
                }
                mid = *lower.begin();
            }
        }
    }
    return 0;
}

三種解法對比

總結

PAT 1057 經典展示了「動態中位數維護」的三種經典數據結構：

1. 離線/小範圍 → 樹狀數組 / 線段樹 (最快，需座標壓縮或固定範圍)。
2. 線上/通用型別 → 雙堆 / 雙平衡 BST (multiset, policy-based data structure)。
3. 面試/工程 → 雙 priority_queue (僅支援中位數) 或 multiset (支援任意順序統計)。

掌握三者權衡，是數據結構功力的重要標誌。