【ybtoj高效进阶 21257】最小路费（换根DP）

最小路费

题目链接：ybtoj高效进阶 21257

题目大意

给你一个树，点有权重，边有边权，然后要你选一个点作为目标点，使得费用最小。
然后费用是每个点到目标点的距离的平方乘它的权重的和。

思路

不难看出这道题其实就是一个换根 DP。

就维护 \(zero_i,one_i,two_i\) 分别表示 \(i\) 的子树内所有点到 \(i\) 的零次项，一次项和二次项。
然后一个点的答案就是它作为根节点的二次项，然后换根 DP 每个点都作为根节点取一个 \(\min\) 就好了。

代码

#include<cstdio>
#include<iostream>
#define ll long long

using namespace std;

struct node {
	ll x;
	int to, nxt;
}e[400001];
int n, c[200001], x, y;
int le[200001], KK;
ll z, zero[200001], one[200001], two[200001], ans;

void add(int x, int y, ll z) {
	e[++KK] = (node){z, y, le[x]}; le[x] = KK;
	e[++KK] = (node){z, x, le[y]}; le[y] = KK;
}

void dfs_1(int now, int father) {//先求出 1 为根的结果
	zero[now] = c[now];
	for (int i = le[now]; i; i = e[i].nxt)
		if (e[i].to != father) {
			dfs_1(e[i].to, now);
			zero[now] += zero[e[i].to];
			one[now] += one[e[i].to] + zero[e[i].to] * e[i].x;
			two[now] += two[e[i].to] + 2 * one[e[i].to] * e[i].x + e[i].x * e[i].x * zero[e[i].to];
		}
}

void dfs(int now, int father) {//换根 DP
	ans = min(ans, two[now]);
	
	for (int i = le[now]; i; i = e[i].nxt)
		if (e[i].to != father) {
			zero[now] -= zero[e[i].to];
			one[now] -= one[e[i].to] + zero[e[i].to] * e[i].x;
			two[now] -= two[e[i].to] + 2 * one[e[i].to] * e[i].x + e[i].x * e[i].x * zero[e[i].to];
			zero[e[i].to] += zero[now];
			one[e[i].to] += one[now] + zero[now] * e[i].x;
			two[e[i].to] += two[now] + 2 * one[now] * e[i].x + e[i].x * e[i].x * zero[now];
			
			dfs(e[i].to, now);
			
			zero[e[i].to] -= zero[now];
			one[e[i].to] -= one[now] + zero[now] * e[i].x;
			two[e[i].to] -= two[now] + 2 * one[now] * e[i].x + e[i].x * e[i].x * zero[now];
			zero[now] += zero[e[i].to];
			one[now] += one[e[i].to] + zero[e[i].to] * e[i].x;
			two[now] += two[e[i].to] + 2 * one[e[i].to] * e[i].x + e[i].x * e[i].x * zero[e[i].to];
		}
}

int main() {
//	freopen("fare.in", "r", stdin);
//	freopen("fare.out", "w", stdout);
	
	scanf("%d", &n);
	for (int i = 1; i <= n; i++) scanf("%d", &c[i]);
	for (int i = 1; i < n; i++) {
		scanf("%d %d %lld", &x, &y, &z);
		add(x, y, z);
	}
	
	dfs_1(1, 0);
	
	ans = two[1];
	dfs(1, 0);
	
	printf("%lld", ans);
	
	return 0;
}