kmjp's blog

競技プログラミング参加記です

AtCoder ARC #030 : D - グラフではない

ARC

勉強になりました。
http://arc030.contest.atcoder.jp/tasks/arc030_4

問題

N要素の整数列X[i]がある。
ここに対するQ個のクエリを処理せよ。

  1. a,b,vが与えられるので、X[a]～X[b]にvを加える。
  2. a,b,c,dが与えられるので、X[c]～X[d]の内容をX[a]～X[b]に上書きする。
  3. a,bが与えられるので、X[a]～X[b]の和を答える。

解法

１つ目と３つ目のクエリはBITを２つ使うテクで解けるが、問題は２つ目。
ここではEditorialに則り、永続平衡二分木を書く。

永続は置いておいて、まず平衡二分木を作る。
その際は以下を参考にRBSTを実装した。
プログラミングコンテストでのデータ構造 2 ～平衡二分探索木編～

さらに永続化するため、クエリ１によるサブツリーの加算値を各ノードに保持しておく。
木の構造が変わるmerge/splitの処理においてこの加算値を遅延評価で伝搬させる。

今回、クエリ２により木のコピーが大量発生する。
処理時間の観点でもメモリ使用量の観点でもコピーを抑えるため、遅延加算値のないノードは同じメモリを共有すると良い。
また、木に変更を加える場合、共有しているノードのデータを破壊しないようその時だけコピーすると良い。

これでも結構な数の木のコピーが生成される。
ただ、必要なノードは最大N個で、その他は未使用のノードなので、未使用領域を回収して使いまわすと良い。
以下のコードでは、一定量コピー量が増えたら一度X[i]の値を再計算した後にすべての木を破棄し、また木を作り直すことで回収した。

struct RBST* getent(struct RBST*);
RBST* deepcopy(RBST* n);
struct RBST {
	RBST *l,*r;
	ll v,sum,toadd;
	int cnt,rank;
	
	static int count(RBST* n){ return n?n->cnt:0;}
	static RBST* merge(RBST *ln,RBST *rn) {
		if(ln==NULL) return rn;
		if(rn==NULL) return ln;
		RBST* nt;
		if(rand()%(ln->cnt+rn->cnt) < ln->cnt) {
			ln->update();
			nt=deepcopy(ln);
			nt->r = merge(nt->r,rn);
		}
		else {
			rn->update();
			nt=deepcopy(rn);
			nt->l = merge(ln, nt->l);
		}
		return nt->update();
	}
	static RBST* merge(pair<RBST*,RBST*> p) { return merge(p.first,p.second);}
	
	pair<RBST*, RBST*> split(int k) {
		update();
		RBST* nt=deepcopy(this);
		if(k<=count(nt->l)) {
			pair<RBST*, RBST*> p=nt->l?nt->l->split(k):make_pair((RBST*)NULL,(RBST*)NULL);
			nt->l=p.second;
			return make_pair(p.first,nt->update());
		}
		else {
			pair<RBST*, RBST*> p=nt->r?nt->r->split(k-count(nt->l)-1):make_pair((RBST*)NULL,(RBST*)NULL);
			nt->r=p.first;
			return make_pair(nt->update(),p.second);
		}
	}
	
	RBST* update() {
		cnt=1;
		rank=0;
		sum=v;
		if(l) cnt+=count(l), sum+=l->sum, rank=max(rank,1+l->rank);
		if(r) cnt+=count(r), sum+=r->sum, rank=max(rank,1+r->rank);
		sum+=toadd*cnt;
		return this;
	}
	
	static RBST* deepcopy(RBST* n) {
		if(n==NULL) return NULL;
		
		n=getent(n);
		if(n->toadd) {
			if(n->l) {
				n->l=getent(n->l);
				n->l->toadd += n->toadd;
				n->l->update();
			}
			if(n->r) {
				n->r=getent(n->r);
				n->r->toadd += n->toadd;
				n->r->update();
			}
			n->v += n->toadd;
			n->toadd = 0;
		}
		return n->update();
	}
	
	RBST* add(int l,int r,ll v) {
		l=max(l,0);
		r=min(cnt,r);
		if(r-l<=0) return this;
		
		RBST* cl;
		if(l==0 && cnt==r) {
			cl = getent(this);
			cl->toadd += v;
		}
		else {
			cl = deepcopy(this);
			if(cl->l && l<cl->l->cnt) cl->l = cl->l->add(l,r,v);
			if(cl->r && count(cl->l)+1<r) cl->r = cl->r->add(l-(count(cl->l)+1),r-(count(cl->l)+1),v);
			if(l<=count(cl->l) && count(cl->l)+1<=r) cl->v+=v;
		}
		return cl->update();
	}
	ll getsum(int k) {
		if(k<=0) return 0;
		if(cnt<=k) return sum;
		
		ll ret=toadd*k;
		
		if(l) ret+=l->getsum(min(k,l->cnt)), k-=min(k,l->cnt);
		if(k) ret+=v, k--;
		if(k && r) ret+=r->getsum(min(k,r->cnt)), k-=min(k,r->cnt);
		assert(k==0);
		return ret;
	}
};

const int resv=15000000;
vector<RBST> mem;
RBST *top,*cur;
int N,Q;
ll X[300000];

RBST *getent(RBST* n){
	if(n==NULL) return NULL;
	mem.push_back(*n);
	return &mem.back();
}

void rebuild() {
	int i;
	
	mem.resize(N);
	top=NULL;
	FOR(i,N) {
		cur=&mem[i];
		cur->l=cur->r=NULL;
		cur->toadd=cur->sum=0;
		cur->v=X[i];
		cur->update();
		top=RBST::merge(top,cur);
	}
}

#include <sys/types.h>
#include <unistd.h>

void solve() {
	int i,j,k,l,r,x,y,a,b,c,d,v; string s;
	
	srand(time(NULL)+getpid());
	cin>>N>>Q;
	mem.reserve(resv);
	
	FOR(i,N) cin>>X[i];
	rebuild();
	
	pair<RBST*,RBST*> p1,p2;
	while(Q--) {
		cin>>i>>a>>b;
		
		if(i==1) {
			cin>>v;
			top=top->add(a-1,b,v);
		}
		else if(i==3) {
			cout<<top->getsum(b)-top->getsum(a-1)<<endl;
		}
		else if(i==2) {
			cin>>c>>d;
			p2=top->split(d);
			p1=p2.first->split(c-1);
			cur=RBST::deepcopy(p1.second);
			top=RBST::merge(getent(p1.first),getent(p1.second));
			top=RBST::merge(getent(top),getent(p2.second));
			
			p2=top->split(b);
			p1=p2.first->split(a-1);
			top=RBST::merge(getent(p1.first),cur);
			top=RBST::merge(getent(top),getent(p2.second));
		}
		// GC?
		if(mem.size()>resv*9/10) {
			FOR(i,N) X[i]=top->getsum(i+1)-top->getsum(i);
			rebuild();
		}
	}
}

まとめ

平衡二分木のいい勉強になりました。
とはいえこれ本番に通すのはつらいな…。