CSAcademy Round #31 : E. Uniform Trees

Fは微妙にTLEで解けきれず。
https://csacademy.com/contest/round-31/task/uniform-trees/

問題

各頂点にラベルが振られた根付き木がある。
Uniformな木とは、ある頂点の子頂点のラベルがすべて同じ場合のものをいう。

この木の頂点の部分集合がUniformであるとは、以下の条件を満たす場合をいう。

• 部分集合中の全頂点のうちある1点は他のすべての頂点の祖先になる。この点を部分集合からなるグラフの根とする。
• 部分集合中の根を除く各頂点は、最寄の祖先と辺をつなぐ。
• こうして作った根付き木がUniformである。

元の根付き木に対し、Uniformな頂点の部分集合を求めよ。

解法

マージテクの要領で、各頂点vのSubTree中から、各頂点対のLCAがすべてvになる範囲で同一ラベルの頂点をいくつか選んだ時の組み合わせを計算しよう。
vのSubtree中の部分集合にvを含める場合、他の頂点は同一ラベルの範囲でしか頂点を選べないので、全ラベルの組み合わせの総和となる。

int N;
int P[202020],V[202020];
vector<int> E[202020];

ll mo=1000000007;
ll tot;

map<int,ll> M[202020];
ll sum[202020];

void dfs(int cur) {
	
	FORR(e,E[cur]) {
		dfs(e);
		if(M[cur].size()<M[e].size()) {
			swap(M[cur],M[e]);
			swap(sum[cur],sum[e]);
		}
		
		FORR(m,M[e]) {
			if(M[cur].count(m.first)) {
				(sum[cur] += M[cur][m.first]*(mo+m.second-1)) %= mo;
				M[cur][m.first] = M[cur][m.first]*m.second % mo;
			}
			else {
				M[cur][m.first] = m.second;
				(sum[cur] += m.second) %= mo;
			}
		}
	}
	
	ll pat = (1+sum[cur]+mo-M[cur].size())%mo;
	(tot += pat)%=mo;
	if(M[cur].count(V[cur])) {
		(sum[cur] += 0+pat)%=mo;
		(M[cur][V[cur]] += 0+pat)%=mo;
	}
	else {
		(sum[cur] += 1+pat)%=mo;
		(M[cur][V[cur]] += 1+pat)%=mo;
	}
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>P[i]>>V[i];
		if(i) {
			P[i]--;
			E[P[i]].push_back(i);
		}
	}
	
	dfs(0);
	
	cout<<tot<<endl;
}

まとめ

手間取ったけど何とか解けてよかった。
少しマージテクの苦手意識が減ってきた気がする。