[oiclass3904]种树：树形DP+换根

题目

题面
BLUESKY007喜欢种树。一天，她得到了一棵\(n\)个点的树，其中节点\(i\)重量为\(w_i\) 。在种树之前，BLUESKY007需要用起重机把树吊起。由于她只有一台起重机，所以她只能选择一个点作为受力点。根据BLUESKY007所在世界的物理知识，吊起一棵树需要做的功为\(\sum_{i=1}^{n}w_i\cdot dis_i\)，其中\(dis_i\)表示节点\(i\)与受力点之间的距离（边数）。

由于吊起这棵树的费用与所做的功正相关，所以BLUESKY007希望所做的功尽可能小。请你帮助她求出吊起这棵树所做的功的最小值。

题解

本题是裸的换根DP，很好实现。设以1为根结点，dep[u]表示u到根的路径长度，size[u]表示以u为根的子树中所有结点的重量和，当以1为根结点时，dfs一遍求出总做功的值，即\(sum=\sum w[u]*dep[u]\)。接着考虑换根带来的影响，当根从u换到儿子结点v时，v所在的子树少做功size[v]，v外的结点多做功size[1]-size[v]，即\(sum'=sum-size[v]+(size[1]-size[v])=sum-2*size[v]+size[1]\)。
要使\(sum'\leq sum\)，则\(sum-2*size[v]+size[1]\leq sum\)，即\(size[1]\leq 2*size[v]\)，所以在换根时，只有满足\(2*size[v]\geq size[1]\)的儿子结点v才可能使得总做功值变小，可以利用这个条件进行剪枝。

代码

点击查看代码

#include<bits/stdc++.h>
using namespace std;
const int N=2*1e5+5;
int n,w[N];
long long size[N],dep[N],sum,ans=1e18;
vector<int> g[N];
void dfs(int u,int fa){
	sum+=w[u]*dep[u];
	size[u]=w[u];
	for(int i=0;i<g[u].size();i++){
		int v=g[u][i];
		if(v==fa)continue;
		dep[v]=dep[u]+1;
		dfs(v,u);
		size[u]+=size[v];
	}
}
void change_root(int u,int fa,long long sum){
	ans=min(ans,sum);
	for(int i=0;i<g[u].size();i++){
		int v=g[u][i];
		if(v==fa)continue;
		long long tmp=sum-2*size[v]+size[1];
		if(2*size[v]>=size[1])change_root(v,u,tmp);
	}
}
int main(){
	scanf("%d",&n);
	for(int i=1;i<=n;i++){
		scanf("%d",&w[i]);
	}
	for(int i=1,u,v;i<n;i++){
		scanf("%d %d",&u,&v);
		g[u].push_back(v);
		g[v].push_back(u);
	}
	dfs(1,-1);
	change_root(1,-1,sum);
	printf("%lld",ans);
}