HourRank18 : D. Find Three Hackers in a Tree

Cで手こずりすぎなければ、満点はともかくTシャツは手に入ったかもしれないな…。
https://www.hackerrank.com/contests/hourrank-18/challenges/find-three-hackers-in-a-tree

問題

N頂点からなる木を成す無向グラフと、3値A,B,Cが与えられる。
３頂点の組(u,v,w)のうち、３つのパスu-v,v-w,w-uのそれぞれの長さを並べ替えるとA:B:Cとなるものがいくつ存在するか求めよ。

解法

O(N^3)解法は容易だがTLEする。
さて、u-v、v-w、w-uの３つのパスを考えると、共通部分は1つの頂点にしかならない。
これをxとする。

xを総当たりしたとき、xを根とした根付き木を考えてみる。
異なる子頂点のSubTree内から、(-A+B+C):(A-B+C):(A+B-C)の比率になるような3頂点u,v,wを選べば、それらは条件を満たす。
よって各頂点の先にxからの距離がいくつの頂点がいくつあるかを求めていき、k*(-A+B+C)、k*(A-B+C)、k*(A+B-C)の距離にあるものをBitDP等で３頂点分数え上げていけばよい。

int N,A[3],B[3];
vector<int> E[2020];


void dfs(int cur,int pre,int d,map<int,int>& M) {
	M[d]++;
	FORR(e,E[cur]) if(e!=pre) dfs(e,cur,d+1,M);
}

void solve() {
	int i,j,k,l,r,x,y,z; string s;
	
	cin>>N>>A[0]>>A[1]>>A[2];
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	B[0]=-A[0]+A[1]+A[2];
	B[1]=A[0]-A[1]+A[2];
	B[2]=A[0]+A[1]-A[2];
	int g=__gcd(B[0],__gcd(B[1],B[2]));
	B[0]/=g;
	B[1]/=g;
	B[2]/=g;
	sort(B,B+3);
	
	ll ret=0;
	FOR(i,N) {
		ll from[1010][8]={};
		ll to[1010][8]={};
		for(k=1;k*(B[0]+B[1]+B[2])<N;k++) from[k][B[0]==0]=to[k][B[0]==0]=1;
		FORR(e,E[i]) {
			map<int,int> M;
			dfs(e,i,1,M);
			for(k=1;k*(B[0]+B[1]+B[2])<N;k++) {
				FOR(x,8) {
					if((x&1)==0                        ) to[k][x^1] += from[k][x]*M[k*B[0]];
					if((x&2)==0 && (B[0]<B[1] || (x&1))) to[k][x^2] += from[k][x]*M[k*B[1]];
					if((x&4)==0 && (B[1]<B[2] || (x&2))) to[k][x^4] += from[k][x]*M[k*B[2]];
				}
			}
			memmove(from,to,sizeof(from));
		}
		for(k=1;k*(B[0]+B[1]+B[2])<N;k++) ret += from[k][7];
	}
	cout<<ret<<endl;
	
}

まとめ

「u-v、v-w、w-uの３つのパスを考えると、共通部分は1つの頂点にしかならない」という性質を覚えておいた方がよさそうね。