Hatena::Grouptopcoder

kuuso1の日記

2015-03-02

Codeforces Round #294 (Div. 2)E. A and B and Lecture Rooms

| 02:03 |  Codeforces Round #294 (Div. 2)E. A and B and Lecture Rooms - kuuso1の日記 を含むブックマーク はてなブックマーク -  Codeforces Round #294 (Div. 2)E. A and B and Lecture Rooms - kuuso1の日記

http://codeforces.com/contest/519/problem/E

  • N 個の教室が N-1 個の廊下で連結されている校舎を考える (つまり木をなしている)
  • A君とB君はそれぞれある教室 x, y でコンテストの問題を解いたのち、合流して答え合わせをする。
  • A君とB君の合流までの移動距離が等しくなるような教室の数を答えよ。
  • なお毎日このように問題を解くため、M日間分について答えよ

(制約)1≤ N ≤10^5, 1≤ M ≤10^5

・クエリが 10^5 と多いため、都度O(N) で探しているとTLEしてしまう。

・中点となる教室を基点として、そこより遠い教室が該当するので、まず中点を求める。木であることを利用してLCAを用いることでO(LogN)で距離は求まる。

・下図のように、深い方から距離の半分上昇すれば中点が求まる。距離が奇数の場合は該当する教室は存在しない。

f:id:kuuso1:20150303014803p:image

・片方が中点の子ノードであることは明らかだが、もう片方が中点の子ノードかどうかは、最初書いたコードでは上図のようにLCAで判定したが、

 絵を描いてみるとdepthを見るだけでもわかる。

・各ノードにぶら下がる子ノードの数や、LCAのためのdepth等はdfsで求めておく。親側のノード数は、全体Nから自分と子ノードの数を引けば分かる。

・コーナーケースは最初から同じ教室にいる場合。d/2 - 1 が負になるのでバグる。

using System;
using System.Collections;
using System.Collections.Generic;

class TEST{
	static void Main(){
		Sol mySol =new Sol();
		mySol.Solve();
	}
}

class Sol{
	public void Solve(){
		parent=new int[lmax][];
		for(int i=0;i<lmax;i++){
			parent[i]=new int[N];
		}
		depth=new int[N];
		
		root=0;
		dfs(root,-1,0);//深さを数えるdfs
		
		for(int i=0;i<lmax-1;i++){
			for(int j=0;j<N;j++){
				if(parent[i][j]<0){
					parent[i+1][j]=-1;
				}else{
					parent[i+1][j]=parent[i][parent[i][j]];
				}
			}
		}
		
		child=new int[N];
		dfs2(root);//子ノードを数えるdfs
		
		int M=ri();
		for(int m=0;m<M;m++){
			var d=ria();
			int u=d[0]-1,v=d[1]-1;
			int ll=lca(u,v);
			int dd=depth[u]-depth[ll]+depth[v]-depth[ll];	//2点の距離
			if(dd%2==1){
				Console.WriteLine(0);
				continue;
			}
			
			dd/=2;
			
			if(dd==0){	//コーナーケース
				Console.WriteLine(N);
				continue;
			}
			if(depth[u]>depth[v]){int tmp=u;u=v;v=tmp;}
			int trgt=v;
			
			for(int i=0;i<lmax;i++){
				if((((dd-1)>>i)&1)>0){
					trgt=parent[i][trgt];
				}
			}
			
			int c=parent[0][trgt];
			int ans=0;
			if(lca(c,u)==c){	// if(depth[u]==depth[v])でも可
				int trgt2=u;
				for(int i=0;i<lmax;i++){
					if((((dd-1)>>i)&1)>0){
						trgt2=parent[i][trgt2];
					}
				}
				ans=(N-child[c])+(child[c]-child[trgt]-child[trgt2]);
			}else{
				ans=child[c]-child[trgt];
			}
			
			Console.WriteLine(ans);
		}
	}
	
	int lca(int u,int v){
		if(depth[u]>depth[v])return lca(v,u);
		for(int i=0;i<lmax;i++){
			if( (((depth[v]-depth[u])>>i)&1) >0 ){
				v=parent[i][v];
			}
		}
		if(u==v)return u;
		for(int i=lmax-1;i>=0;i--){
			if(parent[i][u]!=parent[i][v]){
				u=parent[i][u];
				v=parent[i][v];
			}
		}
		return parent[0][u];
	}
	
	int N;
	List<int>[] E;
	int[][] parent;
	int[] depth;
	static int lmax=24;
	int root;
	int[] child;
	
	void dfs(int now,int par,int dep){
		parent[0][now]=par;
		depth[now]=dep;
		foreach(int v in E[now]){
			if(v!=par){
				dfs(v,now,dep+1);
			}
		}
	}
	
	int dfs2(int now){
		int sum=1;
		foreach(int v in E[now]){
			if(v!=parent[0][now]){
				sum+=dfs2(v);
			}
		}
		return child[now]=sum;
	}
	
	public Sol(){
		N=ri();
		E=new List<int>[N];
		for(int i=0;i<N;i++)E[i]=new List<int>();
		for(int i=0;i<N-1;i++){
			var d=ria();
			E[d[0]-1].Add(d[1]-1);
			E[d[1]-1].Add(d[0]-1);
		}
	}

	static int ri(){return int.Parse(Console.ReadLine());}
	static int[] ria(){return Array.ConvertAll(Console.ReadLine().Split(' '),e=>int.Parse(e));}
}

木の問題は絵を描いてみると分かることが多いので結構好きです。

あとLCAは爆速でグラフを駆け上がってる感じが好きです。大体バイナリー法はどれもそんな感じで好きですが。