换根dp 学习笔记

例题：P2986。

我们先画出一棵树，比如下面这棵以 $4$ 为根的树。

我们不妨试试它的其他根，比如说把 $5$ 作为根节点：

我们可以通过父亲节点的数据来计算孩子节点的：

我们发现，用红框框起来的部分每个点 $u$ 对答案的贡献都少了 $C_u\times di{s}_{4,5}$，而用蓝框框起来的部分每个点 $u$ 都对答案产生了 $C_u\times di{s}_{4,5}$ 的贡献。

这启示我们什么？
我们记录一个 $siz$。$si{z}_i$ 表示当 $1$ 为根节点时，以 $i$ 为根的子树的数目和。$d$ 表示父节点到子节点的边的权值。
那么我们就有 $f_u=f_{fa}+(si{z}_1-si{z}_u)\times d-si{z}_u\times d=f_{fa}+si{z}_1\times d-2\times si{z}_u\times d$。
先跑一次 dfs 求出 $f_1$，在跑一次 $dfs$ 求出剩余的所有，最后比大小即可。

点击查看代码

#include<bits/stdc++.h>

#define int ll
#define ll long long
#define i128 __int128

#define mem(a,b) memset((a),(b),sizeof(a))
#define m0(a) memset((a),0,sizeof(a))
#define m1(a) memset(a,-1,sizeof(a))
#define lb(x) ((x)&-(x))
#define lc(x) ((x)<<1)
#define rc(x) (((x)<<1)|1)
#define pb(G,x) (G).push_back((x))
#define For(a,b,c) for(int a=(b);a<=(c);a++)
#define Rep(a,b,c) for(int a=(b);a>=(c);a--)
#define in1(a) a=read()
#define in2(a,b) a=read(), b=read()
#define in3(a,b,c) a=read(), b=read(), c=read()
#define fst first 
#define scd second 
#define dbg puts("IAKIOI")

using namespace std;

int read() {
	int x=0,f=1; char c=getchar();
	for(;c<'0'||c>'9';c=getchar()) f=(c=='-'?-1:1); 
	for(;c<='9'&&c>='0';c=getchar()) x=(x<<1)+(x<<3)+(c^48);
	return x*f;
}
void write(int x) { if(x>=10) write(x/10); putchar('0'+x%10); }

const int mod = 998244353;
int qpo(int a,int b) {int res=1; for(;b;b>>=1,a=(a*a)%mod) if(b&1) res=res*a%mod; return res; }
int inv(int a) {return qpo(a,mod-2); }

#define maxn 200050

int n,m;
int val[maxn];
vector<pair<int,int> > G[maxn];

int f[maxn],siz[maxn];

void dfs1(int u,int fa,int dep) {
	f[1]+=val[u]*dep;
	siz[u]=val[u];
	for(auto [v,d]:G[u]) if(v!=fa) {
		dfs1(v,u,dep+d);
		siz[u]+=siz[v];
	}
}

void dfs2(int u,int fa,int d) {
	if(fa!=-1) f[u]=f[fa]+(siz[1]-siz[u])*d-siz[u]*d;
	for(auto [v,d]:G[u]) if(v!=fa) dfs2(v,u,d);
}

void work() {
	in1(n);
	For(i,1,n) in1(val[i]);
	For(i,1,n-1) {
		int u,v,d;
		in3(u,v,d);
		G[u].push_back({v,d});
		G[v].push_back({u,d});
	}
	dfs1(1,-1,0);
	dfs2(1,-1,0);
	int ans=1e18;
	For(i,1,n) ans=min(ans,f[i]);
	cout<<ans;
}

signed main() {
//	ios::sync_with_stdio(false); 
//	cin.tie(0); cout.tie(0);
	int _=1;
//	_=read();
//	cin>>_;
	For(i,1,_) {
		work();
	}
	return 0;
}