Codeforces 85 D. Sum of Medians

問題

集合S={a1,a2,..,ak}に対して、Sのsum of Mediansとは、
Σ[i≦k かつ i mod 5 = 3]a[i]により定義される。

最初集合Sは空である。
次のn個のクエリが与えられるので、それらに答えよ。

• add x 集合Sにxを加える。この操作の時点でSにxは存在しない。
• del x 集合Sからxを削除する。この操作の時点でSにxは含まれている。
• sum 集合Sのsum of Mediansを出力する。

制約条件

n≦10^5
xは正の整数かつx≦10^9

方針

平方分割でやったら意味不明なWAがとれなかった。
多分平方分割でも出来るはず。

Segment Treeを使った解法は以下の通り。
最初にxの値を全部列挙しておき、座標圧縮した後で、Segment Treeを2つ作る。
sum[p][i]に、pに対応する区間のうち、区間の左から数えてmod 5でi番目のものを取り出したときの和を持たせる。
cnt[p]に、pに対応する区間に含まれる数の個数を持たせる。

クエリが来たとき、sumはsum[0][2]を出力すればよい。
add来たら、xに対応する区間木の葉のノードをsum[p][0]=xに書き換え、cnt[p]=1に書き換えて、その頂点から根までを更新する。
delが来たら、xに対応する区間木の葉のノードをsum[p][0]=0に書き換え、cnt[p]=0に書き換えて、その頂点から根までを更新する。

更新は、sum[k][i]=sum[k*2+1][i]+sum[k*2+2][iを左の区間に含まれる数mod 5だけずらした値]とすればよい。

ソースコード

const int N=1<<17;
int n;
ll sum[2*N-1][5];
int cnt[2*N-1];

vector<pi> qs;
void update(int k){
	cnt[k]=cnt[k*2+1]+cnt[k*2+2];
	rep(i,5)sum[k][i]=sum[k*2+1][i]+sum[k*2+2][(i-cnt[k*2+1]%5+5)%5];
}
void run()
{
	cin>>n;
	map<int,int> M; int k=0;
	rep(i,n){
		int x; string q; cin>>q;
		if(q=="sum")qs.pb(mp(2,0));
		else{
			cin>>x; qs.pb(mp(q=="add",x));
			M[x]=0;
		}
	}
	fr(i,M)i->second=k++;
	
	rep(i,n){
		int t=qs[i].first, x=qs[i].second;
		if(t==2)cout<<sum[0][2]<<endl;
		else{
			int p=M[x]+N-1;
			cnt[p]=t;
			sum[p][0]=t*x;
			for(p=(p-1)/2;;p=(p-1)/2){
				update(p);
				if(p==0)break;
			}
		}
	}
}