先考虑一个\(O(N^2)\)做法。
设选的两个点为\(x,y\)，则一定可以将树分成两个集合\(A,B\)，使得\(A\)集合所有点都去\(x\)，\(B\)集合所有点都去\(y\)，而这两个集合的分界点就是树上的一条边。于是考虑枚举断哪条边，然后对两边分别跑一遍带权树的重心，统计答案加起来取最小值就行了。
现在进行优化，求树的重心的方法可以参考。
以\(1\)为根建树，\(f[u]\)表示所有点到\(u\)的总距离。（人数乘以距离）转移方程是：
\[f[v]=f[u]+size[1]-size[v]-size[v]\] 可以发现，只有当\(size[v]*2>size[1]\)时\(v\)比\(u\)更优，而且满足\(size[v]*2>size[1]\)的\(v\)数量\(<=1\)。所以我们可以预处理出每个点的子树大小最大的儿子和次大的儿子，每次断边时自下而上修改其所有祖先的\(size\)大小，这时最大儿子可能变小，进而被次大儿子替代，直接判断一下然后走此时的大儿子就行。易得时间复杂度为\(O(NH)\)，这也就是题中提到树的高度不超过\(100\)的原因吧。
#include#include #define INF 2147483647using namespace std;#define Open(s) freopen(s".in","r",stdin);freopen(s".out","w",stdout);#define Close fclose(stdin);fclose(stdout);int s, w; char ch;inline int read(){ s = 0; ch = getchar(); while(ch < '0' || ch > '9') ch = getchar(); while(ch >= '0' && ch <= '9'){ s = s * 10 + ch - '0'; ch = getchar();} return s;}const int MAXN = 100010;struct Edge{ int next, to;}e[MAXN << 1];int head[MAXN], num, a[MAXN], size[MAXN], f[MAXN], val[MAXN], dep[MAXN], son[MAXN], Sson[MAXN], A, B, n, root, ans = INF, cut;inline void Add(int from, int to){ e[++num].to = to; e[num].next = head[from]; head[from] = num; e[++num].to = from; e[num].next = head[to]; head[to] = num;}int getsize(int u, int fa){ size[u] = a[u]; f[u] = fa; dep[u] = dep[fa] + 1; for(int i = head[u]; i; i = e[i].next) if(e[i].to != fa){ getsize(e[i].to, u); size[u] += size[e[i].to]; val[u] += val[e[i].to] + size[e[i].to]; if(size[e[i].to] > size[son[u]]){ Sson[u] = son[u]; son[u] = e[i].to; } else if(size[e[i].to] > size[Sson[u]]) Sson[u] = e[i].to; }}void getans(int u, int now, int all, int &res){ res = min(res, now); int v = son[u]; if(v == cut || size[Sson[u]] > size[son[u]]) v = Sson[u]; //如果size变化后次大大于最大 if(!v) return; if(size[v] * 2 > all) getans(v, now + all - size[v] - size[v], all, res); //如果size[v]*2<=all就没有继续往下走的意义了，因为此时u一定最优}int solve(int u){ for(int i = head[u]; i; i = e[i].next) if(e[i].to != f[u]){ cut = e[i].to; //断边 A = B = INF; for(int now = u; now; now = f[now]) size[now] -= size[e[i].to]; //自下而上修改其父亲size getans(1, val[1] - val[e[i].to] - dep[e[i].to] * size[e[i].to], size[1], A); getans(e[i].to, val[e[i].to], size[e[i].to], B); //求两个集合的答案 ans = min(ans, A + B); for(int now = u; now; now = f[now]) size[now] += size[e[i].to]; //回溯 solve(e[i].to); }}int main(){ //Open("practice"); n = read(); dep[0] = -1; for(int i = 1; i < n; ++i) Add(read(), read()); for(int i = 1; i <= n; ++i) a[i] = read(); getsize(1, 0); solve(1); printf("%d\n", ans); return 0;}