mayoko’s diary

プロコンとかいろいろ。

SRM 690 div1 med: TreeWalker

問題

TopCoder Statistics - Problem Statement

有向グラフ G について, p(G) は, 「頂点 (v, u) について, v -> u のパスがあるような組み合わせの数」と定義する。

今, 無向グラフで木のグラフが与えられる。このグラフのそれぞれの辺の向きを決め, 有向グラフにする。この決め方は 2^(N-1) 通りあるが, それぞれの場合の p(G) を計算した時の合計を求めよ。

解法

全体的な方針としては, 各頂点 v について, それを LCA とする頂点 (u, w) のパスを列挙して答えを求めるイメージです。(u, w) のパスに p 個の辺が使われていたとすると, 2^(n-1-p) 個のグラフについて (u, w) のパスが通っていることになりますが, これを考えやすくするために, 「(u, w) のパスが通っている確率は 2^(-p)」と考えることにしましょう。これは, mod_inverse(2, MOD) を求めておき, 最後に 2^(n-1) を掛け算すれば復元できるので考えやすいです。

各頂点 v について,
dp[v] = (v 以下の部分木の各頂点について, 「v までエッジを張っている」という条件を満たす頂点個数の期待値) とします。dfs の中の ret がこれです。これがわかれば, v を LCA として (u, w) のパスが通っているような頂点の組の数

  • v <-> (他の頂点)
  • (v の子) <-> (v の子)

を計算するのは簡単です。dfs の中身を説明すると,

  • ll ret = 1
    • v 自体は v に確率 1 でたどり着くので デフォで 1 にしておく
  • ll r = dfs(ch, G)*inv2%MOD;
    • ch の部分木の各頂点が v にたどり着く期待値
  • (ans += ret*r*2) %= MOD;
    • (v の子) <-> (v の子2), v <-> (v の子) でできる組み合わせを列挙
    • 一回 sum を取ってから~とかやるとダブってるので NG(最初これで引っかかった)
  • (ret += r) %= MOD;
    • v への道を追加
const int MOD = 1e9+7;
const int inv2 = (1+MOD)/2;

ll ans;

ll dfs(int v, const vector<vi>& G) {
	ll ret = 1;
	for (int ch : G[v]) {
		ll r = dfs(ch, G)*inv2%MOD;
		(ans += ret*r*2) %= MOD;
		(ret += r) %= MOD;
	}
	return ret;
}

class TreeWalker {
public:
	int calc(int N, int A0, int B, int C, int mod) {
		vector<ll> A(N);
		A[0] = A0;
		for (int i = 1; i <= N-2; i++) 
			A[i] = (B*A[i-1]+C)%mod;
		vector<vi> G(N);
		for (int i = 1; i <= N-1; i++) {
			int j = A[i-1]%i;
			G[j].push_back(i);
		}
		ans = 0;
		dfs(0, G);
		for (int i = 0; i < N-1; i++) 
			(ans *= 2) %= MOD;
		return ans;
	}
};