2020年CCPC威海站 C. Rencontre 树形dp求块类任意两点距离和

传送门：https://codeforces.ml/gym/102798/problem/C

题意

$给一棵树，树上的点分为三种。$
$首先分别在三种点中随机选择一点a,b,c，然后找到一个合适的v，使得$
$f=min(dis(a,v)+dis(b,v)+dis(c,v))$

$求出f的期望。$

思路

$当a,b,c确定下来，那么v也会确定下来，因为min，所以我们围绕a,b,c展开。$

$数学上可以推出min(dis(a,v)+dis(b,v)+dis(c,v))=\frac{1}{2}(dis(a,b)+dis(b,c)+dis(c,a))$
$这样我们就不需要考虑v的存在，而怎么求dis呢？$

$显然dis中的点是在树上的，所以总贡献为\frac{{\sum }_{x\in a}{\sum }_{y\in b}dis(x,y)}{|a||b|}+\frac{{\sum }_{x\in c}{\sum }_{y\in b}dis(x,y)}{|c||b|}+\frac{{\sum }_{x\in a}{\sum }_{y\in c}dis(x,y)}{|a||c|}$

$a，b，c的个数输入的时候就知道了，那么就是求所有的dis，即树上任意两点距离和。$

$很多人都知道，树上任意两点之间距离就是dfs一遍树，对于一条边ans+=(n-siz[v])\ast siz[v]\ast w$

$这里同如此，只不过是进化版，是要在三个块中选择不同点求任意距离和。$
$没关系，设siz[u][1/2/3]表示以u为根节点的子树中，1/2/3种类的点有多少个。$

$然后根据siz大小循环转移即可。$

$难点应该就是怎么在不同块中求贡献了。首先dfs1预处理出全部的siz，再dfs2对边求贡献。$

Code

#include "bits/stdc++.h"
using namespace std;

typedef long long ll;
const int N = 1e6 + 10;

struct Edge {
    int v;
    ll w;
};

vector<Edge> g[N];
ll cnt[4], siz[N][4];;
double ans;

void dfs1(int u, int fa) {
    for(auto e : g[u]) {
        int v = e.v;
        if(v == fa) continue;
        dfs1(v, u);
        for(int i = 1;i <= 3; i++) {
            siz[u][i] += siz[v][i];
        }
    }
}

void dfs2(int u, int fa) {
    for(auto e : g[u]) {
        int v = e.v;
        if(v == fa) continue;
        dfs2(v, u);
        for(int i = 1;i <= 3; i++) {
            for(int j = 1;j <= 3; j++) {
                if(i == j) continue;
                ans += 1.0 * ((siz[1][i] - siz[v][i]) * siz[v][j] * 1.0 * e.w / (cnt[i] * cnt[j]) / 2.0);
            }
        }
    }
}

void solve() {
    int n; cin >> n;
    for(int i = 1;i < n; i++) {
        int u, v; ll w; cin >> u >> v >> w;
        g[u].push_back(Edge{v, w});
        g[v].push_back(Edge{u, w});
    }
    for(int i = 1;i <= 3; i++) {
        cin >> cnt[i];
        for(int j = 1;j <= cnt[i]; j++) {
            int x; cin >> x;
            siz[x][i]++;
        }
    }
    dfs1(1, -1);
    dfs2(1, -1);
    cout << fixed << setprecision(12) << ans << endl;
}

signed main() {
    solve();
}