kmjp's blog

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

TopCoder SRM 731 Div1 Medium RndSubTree

手間取ったけどまぁ解けて良かった。
https://community.topcoder.com/stat?c=problem_statement&pm=14837

問題

無限の大きさを持つ二分木がある。
初期状態ですべての頂点は白く塗られている。

ここで以下の処理をk回繰り返す。

  • 現在位置を根頂点に設定する。
  • 現在いる頂点が赤色の場合、左右どちらかの子頂点に等確率で移動、という処理を繰り返す。
  • 最終的に白い頂点に到達したらそこを赤く塗る。

こうすると赤い頂点がk個できる。
全頂点ペア間のパス長の総和の期待値を求めよ。

解法

人によって色々解法がありそう。

頂点群のパス長の総和を求めるテクの1つとして、ある辺の両側に頂点が何個ずつあるかを数える、というものがある。
片方にp個、もう片方にq個あるなら、パス長の総和を求める際その辺はp*q回カウントされる。

これを応用してカウントの期待値を求めよう。
各頂点vについてP(v) := (深さ+1)とする。
根頂点ならP(v)=1で、それ以外はP(v) := P(vの親頂点)+1 である。
深さがP(v)の頂点とP(v)+1の頂点を結ぶ辺は2^P(v)個ある。
よってそのような辺1本に対し、その辺がパス長に何回カウントされうるか確率を求めよう。あとは対称性により同じ深さの頂点間を結ぶ辺2^P(v)は同じカウント回数になるので2^P(v)倍すればよい。

V(x) := P(v)=xとなる頂点のうち、二分木で一番左端の頂点
とする。以下をDPで求めていこう。求めるべきはd(1,1)~d(k,k)なのでO(k^2)で列挙できる。
d(i,x) := i回目の処理の結果、初めてV(x)が赤く塗られた状態になる確率

こなると、(i+1)回目以降初めてV(x+1)の頂点を塗ることができる。
V(x)とV(x+1)の間の辺のカウント回数は、V(x+1)のSubTree内の頂点を塗る回数と、SubTree外の頂点を塗る回数の積である。
残り(k-i)回の処理において、V(x+1)のSubTree内の頂点が塗られる確率は2^(-x)である。

すでにi個の頂点はSubTree外を塗っているので、求める期待値は以下の和である。

  • 残り(k-i)個の処理対象がV(x+1)のSubTree内になった場合、すでに処理したi個の頂点はSubTree外なので \displaystyle d(i,x) \times (k-i) \times 2^{-x} \times iだけ期待値に加算する。
  • 残り(k-i)個の処理対象中2個のペアについて、片方がSubTree内、片方がSubTree外になるとこの辺がカウントされるので、 \displaystyle d(i,x) \times \frac{(k-i)(k-i-1)}{2} \times (2^{-x} \times (1-2^{-x}))だけ期待値に加算する。
ll mo=1000000007;

ll dp[2020][2020];
ll p2[2020];
ll r2[2020];

class RndSubTree {
	public:
	int count(int k) {
		int i,j,x;
		
		p2[0]=r2[0]=1;
		for(i=1;i<=2010;i++) {
			p2[i]=p2[i-1]*2%mo;
			r2[i]=r2[i-1]*((mo+1)/2)%mo;
		}
		
		ZERO(dp);
		dp[0][0]=1;
		ll ret=0;
		for(i=1;i<=k;i++) { // num
			for(x=0;x<=i;x++) {
				(dp[i][x]+=dp[i-1][x]*(mo+1-r2[x]))%=mo;
				(dp[i][x+1]+=dp[i-1][x]*r2[x])%=mo;
				
				// left
				ll add=0;
				(add += (dp[i-1][x]*r2[x])%mo*((((k-i)*(k-i-1)*r2[x+1])%mo*((mo+1-r2[x+1])%mo))%mo))%=mo;
				(add += (dp[i-1][x]*r2[x])%mo*(((k-i)*r2[x+1]%mo*i)%mo))%=mo;
				(ret += add*p2[x+1])%=mo;
			}
		}
		return ret;
	}
}

まとめ

これもコードは短いんだけど妙に手間取った。