SRM 712 Div1 600 AverageVarianceSubtree

問題
ノードに重みがつけられた木が与えられる。
木のすべての部分木について重みの分散を計算し、この分散の平均値を求めよ。

解法
DFSしてノード上でDPする頻出テクニックを使う問題だが、分散を計算するために一工夫必要。すべての部分木の分散を効率よく計算するためには（ルート, 部分木のサイズ）ごとに"重みの二乗和"と"重みの和の二乗"を保持すればよい。

例として、ノード(a, b)からなる木とノード(c, d)からなる木をマージする作業を考えてみる。
木(a, b)の二乗和と木(c, d)の二乗和がわかっていたとすると、マージした木の二乗和はそのままこの二つを足せばいいだけなので簡単。

問題は、和の二乗の方で、
(a + b + c + d)^2 = a^2 + b^2 + c^2 + d^2 + 2ab + 2ac + 2ad + 2bc + 2bd + 2cd
(a + b)^2 = a^2 + 2ab + b^2
(c + d)^2 = c^2 + 2cd + d^2
となるので、単純に足すわけにはいかない。

和の二乗 = 二乗の和 + 2 * 異なる項の積和

となっているので、異なる項の積和をうまく部分問題から計算できればよいことが分かる。

木(a, b)の異なる項の積和 = ab
木(c, d)の異なる項の積和 = cd

木(a, b, c, d)の異なる項の積和 = ab + ac + ad + bc + bd + cd = ab + cd + (a + b)(c + d)

となるので、部分問題の異なる項の積和と部分問題の和を使えば計算できることが分かる。

以上をふまえて、

• sum[v][i] = ノードvを根とするサイズiの木の重み和
• sum_sq[v][i] = ノードvを根とするサイズiの木の重み二乗和
• sum_prd[v][i] = ノードvを根とするサイズiの木の重みの異なる項の積和

を葉から根方向にDPさせればよい。木をマージするときは、マージする相手方のサイズの分だけ自身が繰り返し現れるので、その分を掛けるのを忘れないようにする。

あとこの問題は数値計算の精度が厳しいらしく、計算誤差を小さくするため以下の工夫が必要。

• あらかじめ重みを中心化しておく。
• long doubleを使う。

ソースコード

using namespace std;

#define ALL(x) (x).begin(), (x).end()
#define EACH(itr,c) for(__typeof((c).begin()) itr=(c).begin(); itr!=(c).end(); itr++)  
#define FOR(i,b,e) for (int i=(int)(b); i<(int)(e); i++)
#define MP(x,y) make_pair(x,y)
#define REP(i,n) for(int i=0; i<(int)(n); i++)

int n;
long long sz[55][55];
double long sum[55][55];
double long sum_sq[55][55];
double long sum_prd[55][55];
double long w[55];
vector<int> child[55];

void dfs(int v) {
  REP (i, n+1) {
    sz[v][i] = 0;
    sum[v][i] = sum_sq[v][i] = sum_prd[v][i] = 0.0;
  }

  sz[v][1] = 1;
  sum[v][1] = w[v];
  sum_sq[v][1] = w[v] * w[v];

  for (auto &c : child[v]) {
    dfs(c);
    for (int i = n; i > 1; i--) {
      for (int j = 1; j < i; j++) {
        int k = i - j;
        sz[v][i] += sz[v][j] * sz[c][k];
        sum[v][i] += sum[v][j] * sz[c][k] + sum[c][k] * sz[v][j];
        sum_sq[v][i] += sum_sq[v][j] * sz[c][k] + sum_sq[c][k] * sz[v][j];
        sum_prd[v][i] += sum_prd[v][j] * sz[c][k] + sum_prd[c][k] * sz[v][j] + sum[v][j] * sum[c][k];
      }
    }
  }
}

class AverageVarianceSubtree {
  public:
  double average(vector<int> p, vector<int> weight) {
    n = weight.size();
    double long avg = 0.0;
    REP (i, n) avg += weight[i];
    avg /= n;
    REP (i, n) w[i] = weight[i] - avg;
    REP (i, n) child[i].clear();
    REP (i, p.size()) child[p[i]].push_back(i+1);

    dfs(0);

    long long tot = 0;
    long double ret = 0.0;

    REP (v, n) {
      for (int i = 1; i <= n; i++) {
        tot += sz[v][i];
        ret += sum_sq[v][i] / i - (sum_sq[v][i] + 2 * sum_prd[v][i]) / i / i;
      }
    }
    
    return ret / tot;
  }
};