Yandex-R1 Sum of Medians [线段树]

Yandex Round1的Sum of Medians题意是这样的，对整数有三种操作，插入，删除和求和，求和的话就是求该序列排序后的mod5==3项的和，然后有两点条件，插入之前，该序列中无该元素，删除之前，该序列中有该元素，然后需要求和的时候要求和。

这题有两个地方难想，一个是需要离线处理，一个是用线段树来维护。

先把要使用的整数离散化到[1…n]然后对1…n这个区间建立线段树，维护区间的数字个数及该区间项mod5分别等于0…4的项的和，然后用这两种标记值来递推就可以了，复杂度便是O(NlogN)，代码如下：

#include <iostream>
#include <string>
#include <set>
#include <vector>
#include <map>
using namespace std;

const int MAX = 1e5 + 5;

typedef long long int64;

int n, num[MAX];
int64 cnt[MAX * 10], dp[MAX * 10][5];
char ch[MAX][5];
set<int> S;
map<int, int> m_ii;
vector<int> val;

void ready()
{
    scanf("%d", &n);
    for(int i = 0; i < n; i++)
    {
        scanf("%s", ch[i]);
        if(ch[i][0] != 's')
        {
            scanf("%d", &num[i]);
            S.insert(num[i]);
        }
    }
    //离散化
    val = vector<int> (S.begin(), S.end());
    for(int i = 0; i < val.size(); i++)
        m_ii[val[i]] = i;

    memset(cnt, 0, sizeof(cnt));
    memset(dp, 0, sizeof(dp));
}

void update(int u, int L, int R, int x, int d)
{
    if(L == R)
    {
        cnt[u] += d;
        if(cnt[u])  dp[u][0] = val[x];
        else  dp[u][0] = 0;
        return;
    }
    else
    {
        int M = (L + R) / 2;
        if(x <= M)  update(2 * u + 1, L, M, x, d);
        else  update(2 * u + 2, M + 1, R, x, d);

        cnt[u] = cnt[2 * u + 1] + cnt[2 * u + 2];

        for(int i = 0; i < 5; i++)
            dp[u][i] = dp[2 * u + 1][i] + dp[2 * u + 2][((i - cnt[2 * u + 1]) % 5 + 5) % 5];
    }
}

void out(int u, int L, int R)
{
    printf("cnt[%d] = %d, %d %d %d %d %d\n", u, cnt[u], dp[u][0],
        dp[u][1], dp[u][2], dp[u][3], dp[u][4]);
    if(L == R)  return;
    else
    {
        int M = (L + R) / 2;
        if(L <= M)  out(2 * u + 1, L, M);
        if(M + 1 <= R)  out(2 * u + 2, M + 1, R);
    }
}

void go()
{
    for(int i = 0; i < n; i++)
    {
        if(ch[i][0] == 'a')  update(0, 0, val.size() - 1, m_ii[num[i]], 1);
        else if(ch[i][0] == 'd')  update(0, 0, val.size() - 1, m_ii[num[i]], -1);
        else  printf("%I64d\n", dp[0][2]);
    }
    //out(0, 0, val.size() - 1);
}

int main()
{
    ready();
    go();
}