bzoj 3302 [Shoi2005]树的双中心 树形dp

题面

题目传送门

解法

假设选择的两个点为\(x,y\)

那么，一定存在一条边，使得这条边将这棵树分成离\(x\)最近的和离\(y\)最近的

那么我们枚举这条边，然后在分成的两棵树中找带权重心

假设当前某一棵树的最优解为\(x'\)，考虑它的儿子\(y'\)

如果\(y'\)的权值小于\(x'\)的权值，那么挑最小的一个儿子继续走下去

另一棵树同理

但是这样选出的两个带权重心的划分不一定是这条边

这并没有关系，因为当这两个点作为最优解的时候一定存在正好是这两个点划分的边，而这条边我们一定会计算到，并且这两个点对于这条边就是最优解，所以不影响正确性

因为树的深度较低，所以这样做不会超时

在计算权值的时候，可以先树形dp求出每一个点子树内的权值和子树外的权值，然后计算起来就比较方便了

此题细节比较多，写代码的时候需要注意

时间复杂度：\(O(nh)\)

代码

#include <bits/stdc++.h>
#define int long long
#define N 100010
using namespace std;
template <typename node> void chkmax(node &x, node y) {x = max(x, y);}
template <typename node> void chkmin(node &x, node y) {x = min(x, y);}
template <typename node> void read(node &x) {
	x = 0; int f = 1; char c = getchar();
	while (!isdigit(c)) {if (c == '-') f = -1; c = getchar();}
	while (isdigit(c)) x = x * 10 + c - '0', c = getchar(); x *= f;
}
struct Edge {
	int next, num;
} e[N * 3];
int cnt, ans, tim, a[N], d[N], f[N], g[N], dfn[N], siz[N], sum[N];
void add(int x, int y) {
	e[++cnt] = (Edge) {e[x].next, y};
	e[x].next = cnt;
}
void dfs1(int x, int fa) {
	sum[x] = a[x]; d[x] = d[fa] + 1;
	dfn[x] = ++tim, siz[x] = 1;
	for (int p = e[x].next; p; p = e[p].next) {
		int k = e[p].num;
		if (k == fa) continue;
		dfs1(k, x); f[x] += f[k] + sum[k];
		sum[x] += sum[k], siz[x] += siz[k];
	}
}
void dfs2(int x, int fa) {
	for (int p = e[x].next; p; p = e[p].next) {
		int k = e[p].num;
		if (k == fa) continue;
		g[k] = g[x] + f[x] - (f[k] + sum[k]) + sum[1] - sum[k];
		dfs2(k, x);
	}
}
void dfs(int x, int fa) {
	for (int p = e[x].next; p; p = e[p].next) {
		int k = e[p].num;
		if (k == fa) continue;
		int tx = 1, ty = k, dis = d[k] - d[tx], ss = 0;
		while (true) {
			int mn = g[tx] + f[tx] - f[k] - sum[k] * dis, mni = tx;
			bool fl = 0;
			for (int q = e[tx].next; q; q = e[q].next) {
				int tmp = e[q].num, dd = dis;
				if (d[tmp] < d[tx] || tmp == k) continue; fl = 1;
				if (dfn[tmp] <= dfn[k] && dfn[k] <= dfn[tmp] + siz[tmp] - 1) dd--;
					else dd++;
				int s = g[tmp] + f[tmp] - sum[k] * dd - f[k];
				if (s < mn) mn = s, mni = tmp;
			}
			if (mni == tx || !fl) {ss += mn; break;}
			tx = mni;
			if (dfn[tx] <= dfn[k] && dfn[k] <= dfn[tx] + siz[tx] - 1) dis--;
				else dis++;
		}
		while (true) {
			int mn = f[ty] + g[ty] - (g[k] + (sum[1] - sum[k]) * (d[ty] - d[k])), mni = ty;
			bool fl = 0;
			for (int q = e[ty].next; q; q = e[q].next) {
				int tmp = e[q].num;
				if (d[tmp] < d[ty]) continue; fl = 1;
				int s = f[tmp] + g[tmp] - (g[k] + (sum[1] - sum[k]) * (d[tmp] - d[k]));
				if (s < mn) mn = s, mni = tmp;
			}
			if (mni == ty || !fl) {ss += mn; break;}
			ty = mni;
		}
		chkmin(ans, ss);
		dfs(k, x);
	}
}
main() {
	int n; read(n); cnt = n;
	for (int i = 1; i < n; i++) {
		int x, y; read(x), read(y);
		add(x, y), add(y, x);
	}
	for (int i = 1; i <= n; i++) read(a[i]);
	dfs1(1, 0); dfs2(1, 0); ans = 1ll << 60;
	dfs(1, 0); cout << ans << "\n";
	return 0;
}