PAT1057 Stack（树状数组+倍增）

题目大意

要求维护一个栈，提供压栈、弹栈以及求栈内中位数的操作（当栈内元素\(n\)为偶数时，只是求第\(n/2\)个元素而非中间两数的平均值）。最多操作100000次，压栈的数字\(key\)范围是[1,100000]。

题目分析

前两个操作用\(stack\)就好。

求中位数。暴力做法即使用上优先队列也是稳稳的超时。考虑树状数组。

压栈时，将\(key\)值对应的位置加1。弹栈减1。

求中位数，可以二分求出\(sum[1:p]==(n+1)/2\)最小的\(p\)，即为\(ans\)。复杂度\(O(nlog^2n)\)。

问题已被解决，但是还有进一步优化的空间。

考虑倍增（？）。从高到低枚举\(ans-1\)的每一个二进制位，即求最大的\(p\)使得\(sum[1:p]<(n+1)/2\)。我们知道树状数组\(tree[k]=\sum_{i=k-lowbit(k)+1}^knum[i]\)，也就是说如果我们知道\(\sum_{i=1}^knum[i]=A\)且\((1<<j)<lowbit(k)\)，那么\(\sum_{i=1}^{k+(1<<j)}=A+tree[k+(1<<j)]\)。倍增的时候枚举二进制位的时候，恰巧我们也是从大到小枚举的，满足\(j\)与\(k\)的限制。这样，就将一次树状数组上\(logn\)的查询替换成一次简单的加法。复杂度\(O(nlogn)\)。

#include <bits/stdc++.h>

using namespace std;

int num;

stack<int> st;

int sum[100005];

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

void add(int p, int v) {for (int i = p; i <= 100000; i += lowbit(i)) sum[i] += v;}

/*
int get(int p) {
    int ret = 0;
    for (int i = p; i >= 1; i -= lowbit(i)) ret += sum[i];
    return ret;
}
*/

int main() {
    num = 0;
    while (!st.empty()) st.pop();
    memset(sum, 0, sizeof(sum));

    int n;
    scanf("%d", &n);
    for (int _ = 0; _ < n; ++_) {
        char com[20];
        scanf("%s", com);
        if (strcmp(com, "Push") == 0) {
            int key;
            scanf("%d", &key);
            ++num;
            st.push(key);
            add(key, 1);
        } else if (strcmp(com, "Pop") == 0) {
            if (!st.empty()) {
                int key = st.top();
                printf("%d\n", key);
                --num;
                st.pop();
                add(key, -1);
            } else printf("Invalid\n");
        } else {
            if (!st.empty()) {
                int temp = 0, ans = 0;
                for (int i = 16; i >= 0; --i) {
                    if (ans + (1 << i) > 100000) continue;
                    if (temp + sum[ans + (1 << i)] < (num + 1) / 2) temp += sum[ans += (1 << i)];
                }
                printf("%d\n", ans + 1);
            } else printf("Invalid\n");
        }
    }
    return 0;
}