Codeforces ECR #162 : E. Count Paths

あまり出来の良くなかった回。
https://codeforces.com/contest/1923/problem/E

問題

木を成す無向グラフが与えられる。
各点は色が塗られている。
以下を満たすパスは何通りか。

• 両端の頂点は同じ色である。
• 両端以外の頂点は、両端と異なる色である。

解法

SubTreeにおいて、色ごとに根から到達できるその色の頂点数を数えていこう。
SubTreeを2つマージする際、色ごとにその数の掛け算すれば、条件を満たすパスの数が得られる。
SubTreeごとのカウントは、マージテクで管理していく。

int T,N;
int C[202020];
vector<int> E[202020];
map<int,int> M[202020];
ll ret;

void dfs(int cur,int pre) {
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur);
		ret+=M[e][C[cur]];
		M[e].erase(C[cur]);
		if(M[cur].size()<M[e].size()) swap(M[e],M[cur]);
		FORR2(a,b,M[e]) {
			ret+=1LL*M[cur][a]*b;
			M[cur][a]+=b;
		}
	}
	M[cur][C[cur]]=1;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		ret=0;
		FOR(i,N) {
			cin>>C[i];
			M[i].clear();
			E[i].clear();
		}
		FOR(i,N-1) {
			cin>>x>>y;
			E[x-1].push_back(y-1);
			E[y-1].push_back(x-1);
		}
		dfs(0,0);
		cout<<ret<<endl;
	}
			
		
}

まとめ

こっちは割とシンプル。