浙大 PAT 甲级 1057 Stack 树形数组 二分查找 实现查找中位数功能的栈

求中位数不难，但这题会超时，需要用树形数组+二分查找来求中位数。

一般的思路是用一个真实的栈（STL stack）来控制出栈入栈，用一个数组来记录堆栈中的数，这个数组下标代表某个数，对应的元素值代表该数在堆栈中的个数。求中位数，实际上也就是求第(stack.size()/2)大的数，计算这个数组的部分和就可以得到中位数，每次求中位数的最坏时间复杂度为O(n)，平均时间复杂度为O(n/2)，都是线性。有元素入栈出栈需要对这个数组进行更新，时间复杂度为O(1).

但是，如果数组中记录的不是每个数的个数，也就是说A[i]不再代表等于i的数的个数，而是代表小于i的数的个数。这样可以通过二分法求中位数，时间复杂度降低到O(logn)，但是更新却需要O(n)时间复杂度，因为后面的部分和都要随之改变。

树形数组通过巧妙的数学方法，将求中位数的时间复杂度降低到O(logn)，更新的时间复杂度也降低到O(logn)，是一种很有用的数据结构。我也是遇到这题才知道树形数组，网上找了很多资料，目前发现最简单易懂的就是下面这篇：

http://www.cppblog.com/menjitianya/archive/2015/11/02/212171.html

不仅解释了什么是树形数组，也解释了它的好处，归纳了会用到它的经典场景。

下面贴出AC代码

#include<stdio.h>
#include<stack>
#include<string.h>
#include<memory.h>
#define lowbit(i) ((i)&(-i))
#define maxn 1000001
using namespace std;

int c[maxn];
stack<int> S;

void update(int x, int v)
{
    for (int i = x; i < maxn; i += lowbit(i))
    {
        c[i] += v;
    }
}

int sum(int x)
{
    int ans = 0;
    for (int i = x; i >= 1; i -= lowbit(i))
    {
        ans += c[i];
    }
    return ans;
}

void peekMedian()
{
    int left = 1, right = maxn, mid, k = (S.size() + 1) / 2;
    while (left < right)
    {
        mid = (left + right) / 2;
        if (sum(mid) >= k)
        {
            right = mid;
        }
        else
        {
            left = mid + 1;
        }
    }
    printf("%d\n", left);
}

int main()
{
    memset(c, 0, sizeof(int)*maxn);
    while (S.empty() == false)
    {
        S.pop();
    }
    int N;
    scanf("%d", &N);
    char command[11];
    for (int i = 0; i < N; i++)
    {
        scanf("%s", command);
        if (strcmp(command, "Pop") == 0)
        {
            if (S.empty() == true)
            {
                printf("Invalid\n");
            }
            else
            {
                printf("%d\n", S.top());
                update(S.top(), -1);
                S.pop();
            }
        }
        else if (strcmp(command, "PeekMedian") == 0)
        {
            if (S.empty() == true)
            {
                printf("Invalid\n");
            }
            else
            {
                peekMedian();
            }
        }
        else if (strcmp(command, "Push") == 0)
        {
            int tmp = 0;
            scanf("%d", &tmp);
            S.push(tmp);
            update(tmp, 1);
        }
    }
    return 0;
}