kmjp's blog

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

Codeforces #200 Div1. D. Water Tree

Codeforces

本番終了後、「Dはsegtreeで解けるぞ」という話を聞いてチャレンジ。
結果的に2種類のSegTreeライブラリをそろえる問題に…。
http://codeforces.com/contest/343/problem/D

問題

1番の点を根とする根付き木のグラフが与えられる。
このグラフに対し、以下の3つのいずれかからなるクエリを計Q回処理せよ。

  1. 点vを水で満たす。点vの子孫にあたる点には水が満たされる。
  2. 点vの水を空にする。点vの祖先にあたる点からも水がなくなる。
  3. 点vの水があるかないかを答える。

解法

まず、各点の子孫がどこからどこまでかをDFSで求めておく。
ここで2種類のSegTreeを準備する。

1つめは、1度に範囲に対し値の更新を行い、単一点の値を参照できるSegTree。
2つめは、更新は1点ずつだが、範囲に対し条件を満たす点の値を参照できるSegTree。

前者は水を満たしたクエリ番号を格納し、後者は水を空にしたクエリ番号を格納する。
水を満たすクエリが来たら、1つ目のSegTreeを使い、点vの子孫の範囲にまとめてそのクエリ番号を書き込む。
水を空にするクエリが来たら、2つ目のSegTreeを使い、点vにそのクエリ番号を書き込む。また、各点の最大値を保持するようにSegTreeを更新する。
水の有無を問われたら両SegTreeにおいて点vの水を満たした/空にしたクエリ番号を求め、大きい方を答えればよい。

int N,Q,ID;
int L[500010],R[500010],IDs[500010];
const int NNV=1<<20;
vector<int> E[500010];


// 更新は１点、参照はrangeでcompで比較した値をとる
template<class V,int NV> class SegTree_1 {
public:
	vector<V> val;
	V comp(V l,V r){ return max(l,r);};
	
	SegTree_1(){val.resize(NV*2); clear();};
	void clear() { int i; FOR(i,NV*2) val[i]=0;}
	
	V getval(int x,int y,int l,int r,int k) {
		if(r<=x || y<=l) return 0;
		if(x<=l && r<=y) return val[k];
		return comp(getval(x,y,l,(l+r)/2,k*2),getval(x,y,(l+r)/2,r,k*2+1));
	}
	V getval(int x,int y) { return getval(x,y,0,NV,1);}
	void debug() {
		_P("[%d] ",val[1]);
		for(int i=2;i<NV*2;i++) {
			if((i&(i-1))==0) _P("[%d",val[i]);
			else if((i&(i+1))==0) _P(" %d] ",val[i]);
			else _P(" %d",val[i]);
		}
		_P("\n");
	}
	void update(int entry, V v) {
		entry += NV;
		val[entry]=v;
		while(entry>1) {
			//if(Q==100000 && entry<128) _P("%d ",entry);
			entry>>=1;
			val[entry]=comp(val[entry*2],val[entry*2+1]);
		}
		//if(Q==100000) _P("::\na\n");
	}
};

// 更新はrange、参照は点
template<class V,int NV> class SegTree_2 {
public:
	V nolink;
	vector<V> val;
	SegTree_2(){val.resize(NV*2); clear(); nolink=-1<<30;};
	void clear() { for(int i=0;i<NV*2;i++) val[i]=0; }
	
	V getval(int k) {
		int e=1;
		while(val[e]==nolink) e=e*2+(((k*2)&NV)>0), k*=2;
		return val[e];
	}
	void debug() {
		_P("[%d] ",val[1]);
		for(int i=2;i<NV*2;i++) {
			if((i&(i-1))==0) _P("[%d",val[i]);
			else if((i&(i+1))==0) _P(" %d] ",val[i]);
			else _P(" %d",val[i]);
		}
		_P("\n");
	}
	
	void update(int x,int y,int l,int r, V v,int k) {
		if(l>=r) return;
		if(x<=l && r<=y) 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);
		}
	}
	void update(int x,int y,V v) { update(x,y,0,NV,v,1);}
	
};


int dfs(int cur,int pre) {
	int i;
	IDs[ID]=cur;
	L[cur]=ID++;
	FOR(i,E[cur].size()) {
		int tar=E[cur][i];
		if(tar!=pre) dfs(tar,cur);
	}
	R[cur]=ID;
}

void solve() {
	int f,r,i,j,k,l, x,y;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	dfs(0,-1);
	
	SegTree_2<int,NNV> st_on;
	SegTree_1<int,NNV> st_off;
	
	st_on.update(0,NNV,0);
	FOR(i,NNV) st_off.update(i,1);

	cin>>Q;
	FOR(i,Q) {
		cin>>x>>y;
		y--;
		if(x==1) st_on.update(L[y],R[y],i+2);
		else if(x==2) st_off.update(L[y],i+2);
		else {
			if(st_off.getval(L[y],R[y])>st_on.getval(L[y])) _P("0\n");
			else _P("1\n");
		}
	}
	
	
	return;
}

まとめ

2種類のSegTreeを使い分ける珍しい問題。
SegTree苦手なので、今回2つともライブラリにしておいた。