NJPC 2017 : H - 白黒ツリー

これは本番割とシンプルに掛けたので満足。
http://njpc2017.contest.atcoder.jp/tasks/njpc2017_h

問題

根付き木が与えられる。
各頂点は初期状態で白黒いずれかに塗られている。

この状態で、以下のクエリを順次処理せよ。

• 頂点が指定されるので、そのサブツリー内の頂点の色を白黒反転させる。
• 2頂点が指定されるので、その間の最短経路上の頂点の色が、白黒交互に並ぶか判定する。

解法

後者のクエリの条件を、こういいかえる。

• 後者のクエリの2頂点をU-Vとする。U-Vの最短経路上に、両端の頂点の色が同じ辺(条件を満たさない辺)が1つも存在しないかどうか判定する。

前者のクエリが行われると上記判定がどうなるか考える。
以後の後者のクエリのうち、サブツリー内のみを対象とするものは影響を受けない。
それはサブツリー内全体が白黒反転するためである。
ではどこが影響を受けるかというと、前者のクエリの対象の頂点と、その親頂点を結ぶ辺のみである。
よってサブツリー全体を処理する必要はなく、その1辺のみ処理すればよい。

あとはU-Vの最短経路上に条件を満たさない辺があるかどうかを高速に判定したい。
これはBIT+EulerTour+LCAで高速に求められる。

まずEulerTourで各頂点vに対し訪問順L[v]およびサブツリー内のL[w]の最大値R[v]を求めよう。
頂点vのサブツリーには、訪問順が[L[v],R[v]]の頂点群が格納されることになる。
ここで、各頂点vから根頂点までにある、条件を満たさない辺の数F(v)をBITで管理しよう。

仮に頂点vとその親頂点の間に条件を満たさない辺があるなら、vのサブツリー内の頂点wはF(w)を1増やすことになる。
よってそのようなvに対しBITで[L[v],R[v]]の範囲に1ずつ加算していこう。
条件を満たさない辺に対し上記BITの加算を行うとF(v)が求まる。

あとはクエリに対し、F(U)、F(V)、F(LCA(U,V))の値が一致するなら、U→LCA(U,V)→Vの間に条件を満たさない辺を1つも通らなかったことになる。

int N,Q;
vector<int> E[101010];
int C[101010];
int NG[101010];
int id;
int L[101010],R[101010],rev[101010];

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	V add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
BIT<int,20> bt;

int P[21][200005],D[200005];

void dfs(int cur,int depth) {
	L[cur]=id++;
	rev[L[cur]]=cur;
	D[cur]=depth;
	FORR(e,E[cur]) dfs(e,depth+1);
	R[cur]=id;
}

int lca(int a,int b) {
	int ret=0,i,aa=a,bb=b;
	if(D[aa]>D[bb]) swap(aa,bb);
	for(i=19;i>=0;i--) if(D[bb]-D[aa]>=1<<i) bb=P[i][bb];
	for(i=19;i>=0;i--) if(P[i][aa]!=P[i][bb]) aa=P[i][aa], bb=P[i][bb];
	return (aa==bb)?aa:P[0][aa];               // vertex
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	for(i=1;i<N;i++) {
		cin>>P[0][i];
		P[0][i]--;
		E[P[0][i]].push_back(i);
	}
	
	dfs(0,0);

	FOR(i,N) cin>>C[i];
	for(i=1;i<N;i++) {
		NG[i]=(C[i]==C[P[0][i]]);
		if(NG[i]) {
			bt.add(L[i],1);
			bt.add(R[i],-1);
		}
	}
	FOR(i,19) FOR(x,N) P[i+1][x]=P[i][P[i][x]];
	
	
	cin>>Q;
	while(Q--) {
		cin>>i;
		if(i==1) {
			cin>>x;
			x--;
			if(NG[x]) {
				bt.add(L[x],-1),bt.add(R[x],1);
			}
			else {
				bt.add(L[x],1),bt.add(R[x],-1);
			}
			NG[x]^=1;
		}
		else if(i==2) {
			int u,v;
			cin>>u>>v;
			u--,v--;
			int r=lca(u,v);
			
			if(bt(L[r])==bt(L[u]) && bt(L[r])==bt(L[v])) cout<<"YES"<<endl;
			else cout<<"NO"<<endl;
		}
		
	}
}

まとめ

最初HLDecomposition解が頭に浮かんで面倒だと思ったけど、BIT+LCAに持ち込めたのはよかった。