kmjp's blog

競技プログラミング参加記です

Codeforces #263 Div1 B. Appleman and Tree

典型問題っぽい?
http://codeforces.com/contest/461/problem/B

問題

N頂点からなる木を成すグラフが与えられる。
また、各頂点の色の白黒が与えられる。

このグラフを幾つか辺を切断して複数の連結要素にするとき、1つの連結要素あたり1個だけ黒い頂点が含まれるような分け方は何通りか。

解法

木DPでDFSで解いていく。
各頂点の状態として、

  • 親側にある黒点と子側にある黒点両方と連結可能な場合
  • 親側にある黒点と連結できず、子側にある黒点と連結可能な場合

の2通りが考えられ、それぞれの可能な組み合わせ数を計算していく。

各頂点のDFSにおいて:

  • 対象の頂点が黒の場合:
    • 各子の頂点は親側(=現在の点)と連結することもできるし、孫の黒点と連結することもできるので、各子における「親側にある黒点と子側にある黒点両方と連結可能な場合」を掛け合わせたのが答え。
  • 対象の頂点が白の場合:
    • 親側にある黒点と連結できず、子側にある黒点と連結可能な場合:
      • 頂点は親の黒点と連結できないので、どこかの子の先にある黒点と連結しなければならない。よって、子の1個だけ「親側にある黒点と連結できず、子側にある黒点と連結可能な場合」のケースで、残りの個は「親側にある黒点と子側にある黒点両方と連結可能な場合」のケースを掛け合わせればよい。
      • 下手に計算すると子の数の2乗の時間がかかるが、事前に「親側にある黒点と子側にある黒点両方と連結可能な場合」の積を取っておくと子の数のオーダーで計算できる。
    • 親側にある黒点と子側にある黒点両方と連結可能な場合:
      • 上記いずれかの子の黒点と連結するケースに、頂点が親の黒点と連結し、上記「対象の頂点が黒の場合」と同じ計算方法の組み合わせを加算すればよい。
int N;
int C[100002],K;
vector<int> E[100002];
ll mo=1000000007;
ll memo[1000002][2];

ll rev(ll a) {
	ll r=1, n=mo-2;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int f,i,j,k,l,x,y;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x;
		E[x].push_back(i+1);
	}
	
	FOR(i,N) cin>>C[i];
	for(int cur=N-1;cur>=0;cur--) {
		memo[cur][1]=1;
		FOR(i,E[cur].size()) memo[cur][1]=memo[cur][1]*memo[E[cur][i]][1]%mo;
		if(C[cur]) {
			memo[cur][0]=memo[cur][1];
		}
		else {
			FOR(i,E[cur].size()) memo[cur][0]+=(memo[cur][1]*memo[E[cur][i]][0])%mo*rev(memo[E[cur][i]][1])%mo;
			memo[cur][0]%=mo;
			memo[cur][1]=(memo[cur][1]+memo[cur][0])%mo;
		}
	}
	cout<<memo[0][0]<<endl;
}

まとめ

本番もっとややこしいコードを書いて時間がかかった。