【换根DP】洛谷 P2986 [USACO10MAR] Great Cow Gathering G

题链

https://www.luogu.com.cn/problem/P2986

题解

换根DP教程

换根dp一般分为三个步骤
1、先指定一个根节点
2、一次dfs统计子树内的节点对当前节点的贡献
3、一次dfs统计父亲节点对当前节点的贡献并合并统计最终答案

1. 第一次 dfs，以 1 为根，求出每棵子树的奶牛总数，以及以 1 为聚集点的 ans[1]
2. 设 ans[u] 为以 u 为聚集点的 答案
3. 设 g[v] 为 v 的子树下所有奶牛 聚集到 v 的代价
4. 设 cost 为 u 到孩子节点 v 的代价
5. 现在已知 ans[u]，如何求 ans[v] ?
   1. ans[v] = g[v] + ( ans[u] - (g[v] + cnt[v]*cost) ) + (cnt[1] - cnt[v]) * cost;
   2. 其中 把 等式右边 分为三个部分
   3. 第一个部分：g[v] 为 v 的子树下所有奶牛 聚集到 v 的代价，肯定要加上
   4. 第二个部分：ans[u] - (g[v] + cnt[v]*cost)
      • 其中 g[v] + cnt[v]*cost 为 v 的子树下(包括 v) 所有奶牛到 u 的代价
      • ans[u] - (g[v] + cnt[v]*cost) 的话就是减去 v 的子树(包括 v) 的代价
      • 第一个部分 加 第二个部分 还差 除了 v 子树以外的其他奶牛到 v 的代价，就是第三个部分的代价
   5. 第三个部分：(cnt[1] - cnt[v]) * cost
      • 其中 (cnt[1] - cnt[v]) 表示 除了 v 子树以外的奶牛的个数

于是转移方程也写好了

代码

#include <bits/stdc++.h>
using namespace std;
#define ls rt<<1
#define rs rt<<1|1
#define LL long long
#define MS 100009

int n, m; 
LL cnt[MS];
struct node{
	LL poi, cost;
};
vector<node > vc[MS];
LL ans[MS];

void dfs1(int u, int f, LL sum){
	for(auto nb:vc[u]){
		auto v = nb.poi;
		auto c = nb.cost;
		if(v == f) continue;
		ans[1] += (sum + c) * cnt[v];
		dfs1(v, u, sum + c);
		cnt[u] += cnt[v];
	}
}

void dfs2(int u, int f){
	for(auto nb:vc[u]){
		auto v = nb.poi;
		auto c = nb.cost;
		if(v == f) continue;
		ans[v] = ans[u] + (cnt[1] - cnt[v]*2) * c;
		dfs2(v, u);
	}
}

void solve(){
	cin >> n; 
	for(int i=1;i<=n;i++) cin >> cnt[i];
	for(int i=2;i<=n;i++) {
		int u, v, w; cin >> u >> v >> w;
		vc[u].push_back({v, w});
		vc[v].push_back({u, w});
	}
	dfs1(1, 0, 0);
	dfs2(1, 0);
	LL minn = 9e18;
	for(int i=1;i<=n;i++) minn = min(minn, ans[i]);
	cout << minn << "\n";
}

int main(){
	ios::sync_with_stdio(false);
	int ce = 1;
	//cin >> ce;
	while(ce--) solve();
	
	return 0;
}