PAT 1057 Stack - 樹狀數組維護棧中位數查詢

問題描述

實現帶中位數查詢的棧，支援三種操作（N ≤ 10^5）：

• Push key：入棧
• Pop：出棧並輸出值
• PeekMedian：輸出當前棧中中位數（偶數個取第 N/2 小，奇數取第 (N+1)/2 小）

非法操作（空棧 Pop/PeekMedian）輸出 Invalid。

核心思路：樹狀數組維護頻率

鍵值範圍 [1, 10^5]，用樹狀數組 BIT[i] 維護值 ≤ i 的元素個數。中位數查詢轉化為找第 k 小，利用 BIT 前綴和二分查找。

樹狀數組模板

const int MAXV = 100000;
int bit[MAXV + 5];

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

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

int sum(int idx) {
    int s = 0;
    for (; idx > 0; idx -= lowbit(idx)) s += bit[idx];
    return s;
}

// 找最小的 idx 使得 sum(idx) >= k（即第 k 小）
int findKth(int k) {
    int l = 1, r = MAXV;
    while (l < r) {
        int mid = (l + r) >> 1;
        if (sum(mid) >= k) r = mid;
        else l = mid + 1;
    }
    return l;
}

完整解題代碼

#include 
#include 
#include 
#include 
using namespace std;

const int MAXV = 100000;
int bit[MAXV + 5];
vector stk;

int lowbit(int x) { return x & -x; }
void add(int idx, int delta) {
    for (; idx <= MAXV; idx += lowbit(idx)) bit[idx] += delta;
}
int sum(int idx) {
    int s = 0;
    for (; idx > 0; idx -= lowbit(idx)) s += bit[idx];
    return s;
}
int findKth(int k) {
    int l = 1, r = MAXV;
    while (l < r) {
        int mid = (l + r) >> 1;
        if (sum(mid) >= k) r = mid;
        else l = mid + 1;
    }
    return l;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    int n; scanf("%d", &n);
    char cmd[15];
    while (n--) {
        scanf("%s", cmd);
        if (cmd[1] == 'o') {  // Pop
            if (stk.empty()) puts("Invalid");
            else {
                int val = stk.back();
                add(val, -1);
                stk.pop_back();
                printf("%d\n", val);
            }
        } else if (cmd[1] == 'u') {  // Push
            int x; scanf("%d", &x);
            stk.push_back(x);
            add(x, 1);
        } else {  // PeekMedian
            if (stk.empty()) puts("Invalid");
            else {
                int k = (stk.size() + 1) / 2;
                printf("%d\n", findKth(k));
            }
        }
    }
    return 0;
}

關鍵點分析

1. 數據結構選擇：排序/平衡樹/兩堆法均可，但 BIT 代碼最短、常數最小，適合 10^5 規模
2. 中位數定義：偶數個取第 N/2 小（偏左），對應 k = (size+1)/2
3. 命令解析：用 cmd[1] 區分 Pop/Push/Peek（o/u/e）
4. 輸入輸出優化：關閉 sync、用 scanf/printf，避免 TLE

複雜度分析

總體 O(N log²V)，MAXV=10^5 時極快通過。