yukicoder : No.439 チワワのなる木

cww押しですね。
http://yukicoder.me/problems/no/439

問題

N頂点の木を成す無向グラフがある。
各頂点はcかwの文字が振られている。
異なる３頂点の組(x,y,z)を選択したとき、それが以下の条件を満たすのは何通りか。

• x,y,zの頂点はそれぞれ文字c,w,wが振られている。
• yはx→zに向かう最短経路上にある。

解法

グラフを根付き木とし、DFSで解く。
頂点を訪問したとき、その頂点にwが振られている場合、その頂点をyの候補と考えると、x,zが何通りあるかを考える。
その頂点のある1つの隣接頂点の先のにcがa個あったとし、その隣接頂点の先以外にwがb個あったとすると、条件を満たす頂点の組がa*b個作れる。
上記処理は、サブツリー内のc,wの登場回数を数えていくとO(N)で解ける。

int N;
int TC,TW;
string S;
int CC[101010],CW[101010];
vector<int> E[101010];
ll ret;

void dfs(int cur,int pre) {
	if(S[cur]=='c') CC[cur]=1;
	if(S[cur]=='w') CW[cur]=1;
	
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur);
		CC[cur]+=CC[e];
		CW[cur]+=CW[e];
		if(S[cur]=='w') ret += 1LL*CC[e]*(TW-1-CW[e]);
	}
	
	if(pre!=-1 && S[cur]=='w') ret += 1LL*(TC-CC[cur])*(CW[cur]-1);
	
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>S;
	FORR(c,S) if(c=='c') TC++;
	FORR(c,S) if(c=='w') TW++;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	dfs(0,-1);
	cout<<ret<<endl;
}

まとめ

親方向の頂点の処理が簡単に書けたので良かった。