Codeforces #163 Div2. E. More Queries to Array...

Div2とはいえ、Eになると実装が結構しんどい…
それでは最終問題。
http://codeforces.com/contest/266/problem/E

数列に対してクエリが送られるので、数値を更新したり部分列から計算されるスコアを返す問題。
数列の要素数Nが最大100000個、クエリ数Qも100000個なので、O(NQ)よりも小さく処理しないとダメ。

正式なEditorialよりも、blogのページ中のkraskevich氏のコメントの方がわかりやすかったので、これを見て実装。

基本的には、セグメントツリーを作って数列の部分に対する和を計算しておき、クエリに対してそれらを取り出して返す。
ちょっとややこしいのが、A[L]～A[R]に対して、それぞれ1^k、2^k、3^k、…(R-L+1)^kを掛けて返さないといけないため、セグメントツリーは単純な和で組んではいけない。

ここで、セグメントツリー上で以下の２つの和があるとき、これらを結合することを考える。

• A[L]～A[L+(X-1)]の和
• A[L+X]～A[R]の和

これらを足すためには、後者の値からと係数を多くしなければならない。
ここで、kに対応したA[L+X]～A[R]の和をRight[k]と置くと、以下の様にかける。

• 

これは、A[L+X+i]の係数を(i+1)^kから(X+i+1)^kにしたいわけだが、(X+i+1)^k=X^k+C[k][1] X^(k-1)(i+1) + ... となることを考えるとわかる。
この処理では、より次元が低いときのセグメントツリーの部分和が必要だが、これを次元毎に再帰で処理するとTLEにするため、指定した次元より低い値をまとめてvectorで返す関数を作ったりした。

int N,M;
ll C[6][6];
ll cv[6][1<<18];
ll vt[6][1<<17];
ll vtt[6][1<<17];

int val[1<<18];
ll MO=1000000007;

void update(int cur, int l,int r,int v) {
	int wi=17,x,y,s,j,k;
	ll tmp=0;
	
	if(cur>=1<<17) {
		FOR(j,6) cv[j][cur]=v;
		return;
	}
	y=cur;
	while(y){ wi--; y>>=1;}
	
	s = (cur<<(wi+1)) & ((1<<17)-1);
	//_P("==%d %d %d %d %d %d\n",cur,l,r,s,1<<wi,v);
	
	if(l==s && r==s+(2<<wi)) {
		val[cur]=v;
		FOR(j,6) cv[j][cur]=(v*vtt[j][(2<<wi)])%MO;
		return;
	}
	
	if(val[cur]!=-1) {
		val[cur*2]=val[cur*2+1]=val[cur];
		FOR(j,6) cv[j][cur*2]=cv[j][cur*2+1]=(val[cur]*vtt[j][(1<<wi)])%MO;
		val[cur]=-1;
	}
	
	if(r<=s+(1<<wi)) update(cur*2,l,r,v);
	else if(l>=s+(1<<wi)) update(cur*2+1,l,r,v);
	else if(l<s+(1<<wi) && r>=s+(1<<wi)) {
		update(cur*2,l,s+(1<<wi),v);
		update(cur*2+1,s+(1<<wi),r,v);
	}
	
	FOR(k,6) {
		cv[k][cur]=(cv[k][cur*2]+cv[k][cur*2+1]) % MO;
		for(j=1;j<=k;j++) {
			cv[k][cur]=(cv[k][cur] + ((cv[k-j][cur*2+1]*vt[j][1<<wi])%MO)*C[k][j])%MO;
		}
	}
	
	
}

vector<ll> queryvec(int cur, int l,int r,int k) {
	vector<ll> ret;
	int wi=17,x,y,s,j,i;
	ll tmp=0;
	
	ret.resize(k+1);
	if(cur>=1<<17) {
		FOR(i,k+1) ret[i]=cv[i][cur];
		return ret;
	}
	y=cur;
	while(y){ wi--; y>>=1;}
	
	s = (cur<<(wi+1)) & ((1<<17)-1);
	//_P("?? %d %d %d %d %d %d %lld %lld\n",cur,l,r,s,1<<wi,k,cv[k][cur],cv[0][cur]);
	
	if(l==s && r==s+(2<<wi)) {
		FOR(i,k+1) ret[i]=cv[i][cur];
		return ret;
	}
	if(val[cur]>=0) {
		FOR(i,k+1) ret[i]=(val[cur]*vtt[i][r-l])%MO;
		return ret;
	}
	if(r<=s+(1<<wi)) return queryvec(cur*2,l,r,k);
	if(l>=s+(1<<wi)) return queryvec(cur*2+1,l,r,k);
	
	vector<ll> le=queryvec(cur*2,l,s+(1<<wi),k);
	vector<ll> ri=queryvec(cur*2+1,s+(1<<wi),r,k);
	FOR(i,k+1) {
		ret[i]=(le[i]+ri[i])%MO;
		
		for(j=1;j<=i;j++) {
			ret[i] = (ret[i] + ((C[i][j]*vt[j][s+(1<<wi)-l])%MO)*ri[i-j]) % MO;
		}
	}
	return ret;
}

ll query(int cur, int l,int r,int k) {
	int wi=17,x,y,s,j;
	ll tmp=0;

	if(cur>=1<<17) return cv[k][cur];
	y=cur;
	while(y){ wi--; y>>=1;}
	
	s = (cur<<(wi+1)) & ((1<<17)-1);
	//_P("?? %d %d %d %d %d %d %lld %lld\n",cur,l,r,s,1<<wi,k,cv[k][cur],cv[0][cur]);
	
	if(l==s && r==s+(2<<wi)) return cv[k][cur];
	if(val[cur]>=0) {
		// lからrが同じ値
		return (val[cur]*vtt[k][r-l])%MO;
	}
	if(r<=s+(1<<wi)) return query(cur*2,l,r,k);
	if(l>=s+(1<<wi)) return query(cur*2+1,l,r,k);
	
	vector<ll> ri=queryvec(cur*2+1,s+(1<<wi),r,k);
	tmp = (query(cur*2,l,s+(1<<wi),k)+ ri[k])%MO;
	for(j=1;j<=k;j++) {
		tmp = (tmp + ((C[k][j]*vt[j][s+(1<<wi)-l])%MO)*ri[k-j]) % MO;
	}
	return tmp;
	
}

void solve() {
	int x,y,i,j,res,TM,nc,l,r;
	char str[10];
	
	ll o,p;
	
	GET2(&N,&M);
	
	FOR(i,N) {
		val[i+(1<<17)]=GETi();
		FOR(j,6) cv[j][i+(1<<17)]=val[i+(1<<17)];
	}
	FOR(i,1<<17) {
		vt[0][i]=1;
		for(j=1;j<=5;j++) vt[j][i]=(vt[j-1][i]*i)%MO;
	}
	FOR(j,6) {
		vtt[j][0]=0;
		for(i=1;i<=1<<17;i++) vtt[j][i]=(vtt[j][i-1]+vt[j][i]) % MO;
	}
	
	C[0][0]=1;
	for(x=1;x<=5;x++) {
		C[x][0]=1;
		for(y=1;y<=x;y++) C[x][y]=C[x-1][y-1]+C[x-1][y];
	}
	
	
	//segtree
	for(i=(1<<17)-1;i>=1;i--) {
		x=17;y=i;
		while(y){ x--; y>>=1;}
		val[i]=-1;
		l=i*2;
		r=i*2+1;
		FOR(j,6) {
			cv[j][i]=(cv[j][l]+cv[j][r])%MO;
			for(y=1;y<=j;y++) cv[j][i]=(cv[j][i] + ((C[j][y]*cv[j-y][r])%MO)*vt[y][1<<x]) % MO;
		}
	}
	
	res=nc=0;
	FOR(i,M) {
		GETs(str);
		GET3(&l,&r,&x);
		l--; r--;
		if(str[0]=='=') {
			update(1,l,r+1,x);
		}
		else {
			query(1,l,r+1,x);
			_P("%lld\n",query(1,l,r+1,x));
		}
	}
	
	
	
	return;
}

まとめ

かなりデバッグに手間取った。
この分量を本番で書くのは難しいわ…。
でもあのささいなコメントの意図を理解して最後まで書き切れたのは良かった。