fare 换根dp

Problem F. fare

树形换根dp

换根时，维护新节点的信息，使其所有的信息和当它是根节点时一样即可。
推导一下数学式子，维护一下子树 d^2 的和即可。

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

#define endl '\n'
#define int long long
#define fi first
#define se second
#define pb push_back

#define foa(x, y, z) for(int x = (y), ooo = (z); x <= z; ++x)
#define fos(x, y, z) for(int x = (y), ooo = (z); x >= z; --x)
#define ckmax(x, y) ((x) < (y) ? (x) = (y), 1 : 0)
#define ckmin(x, y) ((x) > (y) ? (x) = (y), 1 : 0)

typedef pair<int, int> pii;
typedef long long ll;
const int mod = 1e9 + 7;
const int inf = 0x3f3f3f3f3f3f3f3f;
const int N = 1e6 + 10;
int n, m;
int a[N];
vector<pii> gr[N];
int f[N], g[N], sd[N], siz[N], sa = 0;
void dfs(int r, int fa, int d) 
{
    f[1] += d * d * a[r];
    siz[r] = a[r];
    for(auto &x : gr[r]) {
        int u = x.fi, w = x.se;
        if(u != fa) {
            dfs(u, r, d + w);
            siz[r] += siz[u];
            sd[r] += siz[u] * w + sd[u]; 
        }
    }
}
void dfs2(int r, int fa) 
{
    for(auto &x : gr[r]) {
        int u = x.fi, w = x.se;
        if(u != fa) {
            int down = - 2 * w *  sd[u] - siz[u] * w * w; // 换根时，该节点下面的差
            int up = 2 * w * (sd[r] - sd[u] - siz[u] * w)  // 上面的变化值
                + (sa - siz[u]) * w * w;
            f[u] = f[r] + down + up;
            sd[u] = sd[r] - siz[u] * w + (sa - siz[u]) * w; // 该点作为根节点时，应确保sd 的准确性
//             printf("%d %d %d %d %d\n", u, r, down, up, f[u]);
            dfs2(u, r);
        }
    }
}
void solve()
{
	cin >> n;
    foa(i, 1, n) {
        cin >> a[i];
        sa += a[i];
    }
    foa(i, 1, n - 1) {
        int x, y, w;
        cin >> x >> y >> w;
        gr[x].pb({y, w});
        gr[y].pb({x, w});
    }
    dfs(1, 0, 0);
    dfs2(1, 0);
    int res = *min_element(f + 1, f + n + 1);
    cout  << res << endl;
//     printf("%d %d %d\n", f[1], f[2], f[3]);
}
signed main()
{
	// ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    
	solve();
    return 0;
}