AtCoder ABC #340 (鹿島建設プログラミングコンテスト2024) : G - Leaf Color

Gはかなり早く解けたのに、Fで手間取ったのが残念。
https://atcoder.jp/contests/abc340/tasks/abc340_g

問題

木を成す無向グラフが与えられる。
各点には色が設定されている。
このグラフの部分誘導グラフのうち、葉となる点の色がすべて一致するのは何通りか。

解法

DFSしながら、各SubTree内で根頂点が葉となる、または2つ以上の子頂点のSubTree内で同色の葉を1個以上選ぶケースを数え上げよう。
f(v,c) := 頂点vのSubTreeのうち、vを含みかつ葉として色cのものだけを持つような部分誘導グラフの組み合わせの個数
として、f(v,c)をマージテクの要領で葉から数え上げて行くと良い。

int N;
int A[202020];
vector<int> E[202020];

ll ret;
const ll mo=998244353;
map<int,ll> M[202020];

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

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>A[i];
		A[i]--;
	}
	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+N)%mo<<endl;
	
}

まとめ

もっだいない…。