mayoko’s diary

プロコンとかいろいろ。

SRM 677 div1 med: DiameterOfRandomTree

解法

辺の重みの付け方は 2^(n-1) 通りしかないので, 雰囲気としては dp[v][d] = (頂点 v を root とする部分木で, 直径は d 以下である場合の数) とすれば, 直径が d である確率は, (dp[0][d] - dp[0][d-1]) / (1<<(n-1)) と書くことが出来ます。

あとは dp を求められれば良いですが, 頂点 v を root とする部分木が直径 d 以下になる条件は,
・v の子 ch を root とする部分木の直径が d 以下
・v を経由する path の長さが d 以下
の 2 条件が成り立っていなければなりません。

1 つ目の条件は今の dp でも表現できそうですが, 2 つ目の条件は, 頂点 v を root とする部分木の, 最大深さの情報も知りたいです。これを知っていれば, 2 つの子の最大深さの組み合わせを調べることによって, 2 つ目の条件も確かめられそうです。

ということで, dp を修正して, dp[v][d][s] = (頂点 v を root とする部分木で, 直径は d 以下, 最大深さは s 以下であるような場合の数) とします。
この場合は, 直径が d である確率は, (dp[0][d][(大きい数)] - dp[0][d-1][(大きい数)]) / (1<<(n-1)) とすれば良いです。大きい数は, 適当に 100 とか取っておけば良いです。

dp の遷移を考えます。dp[v][d][s] を求めたいです。
部分木 v の最大深さが m であり, その深さをなす子が ch であるとします(厳密には, ch としては, 何通りか候補がある可能性があるので, 最大深さが m になる, 最も左側の子を ch とします。)。
すると, ch 以外の子は, 深さはまず d-m に制限されます。そうしないと直径が d を超えるので。
また, ch より左側にある子は, さらに深さが m-1 以下に制限されます。ch より右側にある子は深さが m 以下です。各部分木の直径の制約は相変わらず d 以下で変わらないので, 要するに各子について, dp[c][d][up] を解けば, それを掛け算する感じになります。ただし, up は, c==ch の時は up = m, c が ch より左側にある子なら min(m-1, d-m), 右側にある子なら min(m, d-m) です。

ll dp[55][110][110];
vector<int> G[55];
int n;

void dfs(int v, int p, const vector<int>& a, const vector<int>& b) {
    for (int i = 0; i < n-1; i++) {
        if (v == a[i]) {
            if (p == b[i]) continue;
            G[v].push_back(b[i]);
            dfs(b[i], v, a, b);
        }
        if (v == b[i]) {
            if (p == a[i]) continue;
            G[v].push_back(a[i]);
            dfs(a[i], v, a, b);
        }
    }
}

// 頂点 v が root の部分木で, 直径が d 以下, 深さ s 以下であるような場合の数
ll f(int v, int d, int s) {
    ll& ret = dp[v][d][s];
    if (ret >= 0) return ret;
    if (G[v].size() == 0) {
        ret = 1;
        return ret;
    }
    ret = 0;
    int cnt = G[v].size();
    for (int m = 0; m <= d && m <= s; m++) {
        for (int i = 0; i < cnt; i++) {
            for (int w = 1; w <= 2; w++) {
                if (m-w < 0) break;
                ll tmp = f(G[v][i], d, m-w);
                if (m-w-1 >= 0) {
                    tmp -= f(G[v][i], d, m-w-1);
                }
                for (int j = 0; j < cnt; j++) {
                    if (i==j) continue;
                    int up = d-m;
                    if (j < i) {
                        // depth は m 未満
                        up = min(up, m-1);
                    } else {
                        // depth は m 以下
                        up = min(up, m);
                    }
                    ll mul = 0;
                    for (int w2 = 1; w2 <= 2; w2++) {
                        if (up-w2 < 0) break;
                        mul += f(G[v][j], d, up-w2);
                    }
                    tmp *= mul;
                }
                ret += tmp;
            }
        }
    }
    return ret;
}

class DiameterOfRandomTree {
public:
    double getExpectation(vector <int> a, vector <int> b) {
        n = a.size()+1;
        for (int i = 0; i < n; i++) G[i].clear();
        dfs(0, -1, a, b);
        ll prev = 0;
        double ans = 0;
        memset(dp, -1, sizeof(dp));
        for (int i = 0; i <= 2*(n-1); i++) {
            ll c = f(0, i, 2*n);
            ans += i*1.*(c-prev)/(1ll<<(n-1));
            prev = c;
        }
        return ans;
    }
};