kmjp's blog

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

TopCoder SRM 690 Div1 Medium TreeWalker

本番Easyで時間取りすぎてMediumは時間切れ。
https://community.topcoder.com/stat?c=problem_statement&pm=14237

問題

N頂点からなる木を成す無向グラフが与えられる。
(N-1)本の辺が、それぞれどちらか一方の向きの有向辺である場合を考える。
そうしてできる有効グラフの組み合わせは2^(N-1)通りある。
各有向グラフにおいて、u→vと至る経路が存在するような頂点対(u,v)の数の総和を求めよ。

解法

この問題を言い換えると、各辺が等しい確率でどちらかの向きを向く時、移動可能な頂点対の数の期待値(に2^(N-1)を掛けたもの)を答える問題となる。

元の無向グラフを根付き木と見なして考える。
各頂点xに対し、x=LCA(u,v)となるxのsubtree内の2頂点u,vに対し、u→xに到達可能な確率とx→vに到達可能な確率の総和を取ろう。
x→vに到達する確率と、逆にv→xに到達する確率は等しいので、結局u→xに到達可能な確率とv→xに到達可能な確率の総和を取ればよい。

もう少し詳細に書くと、
f(v) := vのsubtree内の頂点群のうち、頂点vまで到達可能な頂点数の期待値
とすると、a,bを異なるxの子頂点としたとき、f(a)/2*f(b)/2の総和を取ればよい。
また、各子頂点a以下の頂点からxに到達する頂点数の期待値として、f(a)の総和も考慮する。

合わせると f(x) = \sum_a \frac{f(a)}{2} + \sum_{a \ne b} \frac{f(a)f(b)}{4}
となる。f(x)の総和に2^(N-1)を掛けたものが解。
(下記のコードでは上記sumの処理でa<bとしている代わりに、最後に2^(N-1)倍する部分を2^N倍して2倍にしている)

int A[101010];
int P[101010];
vector<int> E[101010];
ll dp[101010];
ll mo=1000000007;
int rev2=(mo+1)/2;

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

class TreeWalker {
	public:
	
	int calc(int N, int A0, int B, int C, int MOD) {
		int i;
		
		ZERO(dp);
		
		A[0]=A0;
		FOR(i,N) E[i].clear();
		for(i=1;i<=N-1;i++) {
			A[i] = (1LL*B*A[i-1] + C)%MOD;
			P[i]= A[i-1]%i;
			E[P[i]].push_back(i);
		}
		
		ll ret=0;
		for(i=N-1;i>=0;i--) {
			dp[i]=1;
			FORR(r,E[i]) ret += dp[i]*dp[r]%mo, (dp[i]+=dp[r])%=mo;
			dp[i] = dp[i]*rev2%mo;
			ret %= mo;
		}
		
		return ret*modpow(2,N)%mo;
	}
}

まとめ

以外とコードは短いのよね。