python/TangDou/Topic/HuanGenDp/P1364_HuanGenDp.cpp

#include <bits/stdc++.h>
using namespace std;
const int N = 1e5 + 10, M = N << 1;
const int INF = 0x3f3f3f3f;

// 链式前向星
int e[M], h[N], idx, w[M], ne[M];
void add(int a, int b, int c = 0) {
    e[idx] = b, ne[idx] = h[a], w[idx] = c, h[a] = idx++;
}
int c[N];
int f[N], sz[N];
int ans = INF;

// 第一次dfs,获取在以1为根的树中:
// 1、每个节点分别有多少个子节点，填充sz[]数组
// 2、获取到f[1],f[1]表示在1点设置医院的代价
// 获取到上面这一组+一个数据，才能进行dfs2进行换根
void dfs1(int u, int fa, int step) {
    sz[u] = c[u];
    for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v == fa) continue;
        dfs1(v, u, step + 1);
        sz[u] += sz[v];
    }
    f[1] += step * c[u]; // 先算出1点的总代价
}

void dfs2(int u, int fa) {
    for (int i = h[u]; ~i; i = ne[i]) {
        int v = e[i];
        if (v == fa) continue;
        f[v] = f[u] + sz[1] - sz[v] * 2;
        dfs2(v, u);
    }
    ans = min(ans, f[u]);
}

int main() {
    // 初始化链式前向星
    memset(h, -1, sizeof h);

    int n;
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> c[i];
        int a, b;
        cin >> a >> b;
        if (a) add(a, i), add(i, a); // 是一个二叉树结构，与左右节点相链接，但有可能不存在左或右节点，不存在时,a或b为0
        if (b) add(b, i), add(i, b);
    }
    // 1、准备运作
    dfs1(1, 0, 0);
    // 2、换根dp
    dfs2(1, 0);
    // 输出答案
    cout << ans << endl;
    return 0;
}