Educational DP Contest : V - Subtree

全方位木DPはさほど苦手感ないのであまり戸惑わなかった。
https://atcoder.jp/contests/dp/tasks/dp_v

問題

N頂点の木が与えられる。
そのうち一部の頂点を黒く塗りたい。その際、黒く塗られた頂点は連結であるようにしたい。
各頂点vを含む塗り方はそれぞれ何通りか、mod Mで求めよ。

解法

全方位木DPで解く。
f(v,p) := 根付き木において、頂点vの親がpであるとき、vを黒く塗る場合のvのSubTreeの塗り分け方

とすると、vの子頂点cに対しf(v,p) = prod(1+f(c,v))となる。
また、解はvの隣接点cに対しprod(1+f(c,v))となる。

根が固定されていればf(v,p)は葉から順に埋められるが、根が変わったときはf(p,v)が必要になるケースもある。
ここで全方位木DPを行う。そのためにはまず根を仮決めしてDFSしたのち、再度DFSして親側の組み合わせ数を子に順次伝えていく。
2回目のDFSで子に潜る際、他の子cと親pに対し(1+f(p,v)*prod(1+f(c,v))が必要になるが、prod(1+f(c,v))を毎回求めるとウニ形状でO(N^2)かかりTLEする。
また、Mが素数でないので逆数を掛けることはしたくない。

このような「配列のうち１要素を除いた積を取る」場合は、定番のテクとしてprefixとsuffixの累積積をそれぞれ取っておいて掛けるというテクが使える。

int N;
ll mo;
vector<int> E[101010];
vector<ll> dp[101010];

ll dfs(int cur,int pre) {
	ll ret=1;
	int i;	
	FOR(i,E[cur].size()) {
		int e=E[cur][i];
		if(e!=pre) {
			dp[cur][i]=dfs(e,cur);
			(ret*=(1+dp[cur][i]))%=mo;
		}
	}
	return ret;
}

void dfs2(int cur,int pre,ll p) {
	int i;
	FOR(i,E[cur].size()) {
		int e=E[cur][i];
		if(e==pre) dp[cur][i]=p;
	}
	vector<ll> L(E[cur].size()),R(E[cur].size());
	FOR(i,E[cur].size()) {
		L[i]=(i?L[i-1]:1)*(1+dp[cur][i])%mo;
	}
	for(i=E[cur].size()-1;i>=0;i--) {
		R[i]=((i<E[cur].size()-1)?(R[i+1]):1)*(1+dp[cur][i])%mo;
		if(E[cur][i]!=pre) {
			ll tot=(i?L[i-1]:1)*((i<E[cur].size()-1)?R[i+1]:1)%mo;
			dfs2(E[cur][i],cur,tot);
		}
		
	}
	
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>mo;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	FOR(i,N) dp[i].resize(E[i].size());
	ll tot=dfs(0,-1);
	dfs2(0,-1,0);
	
	FOR(i,N) {
		ll tot=1;
		FORR(d,dp[i]) (tot*=(1+d))%=mo;
		cout<<tot<<endl;
	}
	
}

まとめ

実装は手間だけど苦戦はしなかったかな。