TopCoder SRM 677 Div1 Medium: DiameterOfRandomTree

解法

すぎむさんのコードを参考にしました。木DP。

dp[u][x]:= 頂点uから生えるの最大の辺の長さがxとなる確率

を保存しながらDFSして、作りうる辺の最大値がmax_lを超えないように計算していく。すると、確率の合計が「辺がmax_l以下になる確率」になる。これをP[max_l]とする。P[i]-P[i-1]が直径がiとなる確率なので、期待値が求まる。

コード

#include <cstdio>
#include <cmath>
#include <cstring>
#include <ctime>
#include <iostream>
#include <algorithm>
#include <set>
#include <vector>
#include <sstream>
#include <typeinfo>
#include <fstream>

using namespace std;
typedef long long ll;

class DiameterOfRandomTree {
 public:
  vector<vector<int>> graph;

  void dfs(int u, int p, int max_l, vector<vector<double>> &dp) {
    // dp[u][x]:= 頂点uから生えるの最大の辺の長さがxとなる確率
    dp[u][0] = 1.0;
    for (int v : graph[u]) {
      if (v == p) continue;
      dfs(v, u, max_l, dp);

      vector<double> dp_u(max_l + 1, 0.0);
      for (int x = 0; x <= max_l; ++x)
        for (int dy = 1; dy <= 2; ++dy)
          for (int y = 0; x + y + dy <= max_l; ++y) {
            int z = max(x, y + dy);
            dp_u[z] += dp[u][x] * dp[v][y] * 0.5;
          }

      dp[u].swap(dp_u);
    }
  }

  double getExpectation(vector<int> a, vector<int> b) {
    int N = a.size() + 1;
    int M = 2 * N;
    graph.resize(N);
    for (int i = 0; i < N - 1; ++i) {
      graph[a[i]].push_back(b[i]);
      graph[b[i]].push_back(a[i]);
    }

    vector<double> P(M, 0.0);
    for (int max_l = 0; max_l < M; ++max_l) {
      vector<vector<double>> dp(N, vector<double>(max_l + 1));
      dfs(0, -1, max_l, dp);

      for (int x = 0; x <= max_l; ++x) P[max_l] += dp[0][x];
    }
    double ans = 0.0;
    for (int max_l = 0; max_l < M; ++max_l) {
      ans += max_l * (P[max_l] - P[max_l - 1]);
    }

    return ans;
  }
};