Codeforces #254 Div1 C. DZY Loves Colors

かなり自信なかったけどなんとか本番解けた。
http://codeforces.com/contest/444/problem/C

問題

N要素の数列A[i],B[i]は初期状態でA[i]=i、B[i]=0である。
以下のクエリをM個処理せよ。

• l,r,vが与えられるので、l<=i<=rなiに対しB[i]+=abs(A[i]-v), A[i]=vを行う
• l,rが与えられるので、l<=i<=rなiに対しB[i]の総和を取る

解法

A[i]の現在値とB[i]の総和の２つのデータを管理するSegTreeを作りクエリを処理していく。
1つめのクエリで、SegTreeの領域[x,y]に対応するSegTree上のindexがあって、x<=i<=yにおいてA[i]がすべて同じ値なら、A[id]をvに書き換える際、同時にB[id]+=abs(A[id]-v)*(y+1-x)を処理する。
各クエリO(logN)で処理できるので、全体でO(M*logN)。

template<class V,int NV> class SegTree_2 {
public:
	V nolink;
	vector<V> val;
	vector<V> cur,sum;
	SegTree_2(){val.resize(NV*2); sum.resize(NV*2); cur.resize(NV*2); clear(); nolink=-1<<30;};
	void clear() { for(int i=0;i<NV*2;i++) val[i]=nolink,sum[i]=cur[i]=0; }
	
	V getval(int x,int y,int l,int r,int k) {
		x=max(l,x);
		y=min(r,y);
		if(r<=x || y<=l) return 0;
		if(x<=l && r<=y) return sum[k];
		return getval(x,y,l,(l+r)/2,k*2)+getval(x,y,(l+r)/2,r,k*2+1)+cur[k]*(y-x);
	}
	V getval(int x,int y) { return getval(x,y,0,NV,1);}
	
	void update(int x,int y,int l,int r, V v,int k) {
		if(l>=r) return;
		x=max(l,x);
		y=min(r,y);
		if(x>=y) return;
		if(x<=l && r<=y) {
			if(val[k]!=nolink) {
				cur[k]+=abs(val[k]-v);
			}
			else {
				update(x,y,l,(l+r)/2,v,k*2);
				update(x,y,(l+r)/2,r,v,k*2+1);
			}
			if(k*2<NV*2) sum[k]=sum[k*2]+sum[k*2+1]+cur[k]*(r-l);
			else sum[k]=cur[k]*(r-l);
			val[k]=v;
		}
		else if(l < y && x < r) {
			if(val[k]!=nolink) {
				val[k*2]=val[k*2+1]=val[k];
				val[k]=nolink;
			}
			update(x,y,l,(l+r)/2,v,k*2);
			update(x,y,(l+r)/2,r,v,k*2+1);
			if(k*2<NV*2) sum[k]=sum[k*2]+sum[k*2+1]+cur[k]*(r-l);
			else sum[k]=cur[k]*(r-l);
		}
	}
	void update(int x,int y,V v) { update(x,y,0,NV,v,1);}
};

int N,M;
SegTree_2<ll,1<<20> st;

void solve() {
	int f,i,j,k,l,x,y;
	
	cin>>N>>M;
	FOR(i,N) st.val[i+(1<<20)]=i+1;
	FOR(i,1<<20) st.val[i]=st.nolink;
	
	while(M--) {
		cin>>i>>x>>y;
		if(i==1) {
			cin>>k;
			st.update(x-1,y,k);
		}
		else {
			_P("%lld\n",st.getval(x-1,y));
		}
	}
}

まとめ

だいぶややこしいSegTreeだが、本番に解ききれてよかった。