Codeforces #177 Div1 D. Polo the Penguin and Trees

これも何とか自力回答。
http://codeforces.com/contest/288/problem/D

問題

1～Nの数字が振られた頂点からなる木が与えられる。
1～Nのうちから、a＜b、c＜dとなる４つの数字を選んだ時、a番とb番の頂点の最短路と、c番とd番の頂点の最短路が同じ頂点を通らないような(a,b,c,d)の組は何通りあるか。

解法

aとb、cとdの大小関係は置いておいて、２つのパスが重ならないように取る取り方を考えればよい。

まず１回DFSで、各頂点において各辺の先に何頂点あるか求めておく。

次に各頂点xをa-bの一部の候補として、以下のケースが何通りあるか答える。

• xがa-bの端である場合、xの辺を１個選択するとその先の頂点のいずれかがもう１端であり、その組みわせはその辺の先の頂点数となる。
  • さらに、もう１辺を選び、その辺の先でc-dのパスを構築する。これは辺の先にある頂点群から２つ選ぶ選び方と一致する。
• xがa-bの端でない場合、xの辺を２個選択するとそれぞれの先の頂点のいずれかがもう１端であり、その組みわせはその辺の先の頂点数となる。
  • 上記同様さらに、もう１辺を選び、その辺の先でc-dのパスを構築する。これは辺の先にある頂点群から２つ選ぶ選び方と一致する。

上記処理だと、xを選択し、a-bの端となる辺を２つ選択し、c-dのパスを構築する辺を選択するので最悪O(N^4)かかり間に合わない。
うまく累積和を使うと、全体でO(N)で済むようになる。

int N;
vector<ll> E[100000],D[100000];

int dfs(int cur,int pre) {
	int i,j=-1,k=N-1;
	D[cur].resize(E[cur].size());
	FOR(i,E[cur].size()) {
		int tar=E[cur][i];
		if(tar==pre) j=i;
		else k-=D[cur][i]=dfs(tar,cur);
	}
	if(j>=0) D[cur][j]=k;
	return N-k;

}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	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);
	ll ret=0;
	FOR(i,N) {
		ll tot1=0,tot2=0,tot5=0;
		FOR(j,D[i].size()) tot1+=D[i][j],tot2+=D[i][j]*(D[i][j]-1)/2, tot5+=D[i][j]*D[i][j]*(D[i][j]-1)/2;
		FOR(j,D[i].size()) ret+=D[i][j]*(tot2-D[i][j]*(D[i][j]-1)/2);
		ll tot3=tot1, tot4=tot1;
		FOR(j,D[i].size()) tot3-=D[i][j], ret+=D[i][j]*tot3*tot2;
		FOR(j,D[i].size()) tot4-=D[i][j],  ret-=D[i][j]*tot4*(D[i][j]*(D[i][j]-1)/2);
		FOR(j,D[i].size()) tot5-=D[i][j]*D[i][j]*(D[i][j]-1)/2, ret-=D[i][j]*tot5;
	}
	
	cout << ret << endl;
}

まとめ

なんかこれは割とすぐに思いついた。
むしろ累積和の取り方に苦戦。