首先,我們會有一個很簡單的想法,列舉斷邊,產生兩棵子樹,然後在兩棵樹內分別求帶權重心,計算貢獻,這樣的話複雜度是 \(O(n^2)\) 的。
那麼我們要好好利用 $h \leq 100 $ 的性質。
考慮 \(sze[u]\) 為帶權重量,\(g[u]\) 為以 \(u\) 為根的樹,所有點都到 \(u\) 的代價。
所以 \(g[u] = \sum\limits_{v\in{son(u)}}{g[v] + sze[v]}\)
考慮 \(f[u]\) 為 \(u\) 在的一棵樹中,所有點到 \(u\) 的總代價。
由於計算是動態的,我們考慮轉移。
對於 \(v \in son(u)\) 有 \(f[v] = f[u] + (S - sze[v]) - sze[v]\)
如果 \(v\) 比 \(u\) 更優,當且僅當 \(2 * sze[v] > S\)
我們發現,只有兩個候選項,重兒子和次重兒子。遞迴時順帶維護變化即可。
最壞情況可能會是一條到底的鏈,所以複雜度 \(O(nh)\)
code:
const int Maxn = 5e5 + 5, Maxm = 1e6 + 5;
int n, cnt = 0, w[Maxn], head[Maxn], ver[Maxm], nxt[Maxm];
inline void AddEdge(int u, int v) {
ver[++cnt] = v, nxt[cnt] = head[u], head[u] = cnt;
ver[++cnt] = u, nxt[cnt] = head[v], head[v] = cnt;
}
int dep[Maxn], fat[Maxn], son[Maxn], son2[Maxn];
int sze[Maxn], g[Maxn], cut, ans = INT_MAX;
inline void DfsFir(int u) {
sze[u] = w[u];
for (int i = head[u]; i; i = nxt[i]) {
if (ver[i] == fat[u]) continue;
dep[ver[i]] = dep[u] + 1; fat[ver[i]] = u;
DfsFir(ver[i]); sze[u] += sze[ver[i]];
g[u] += g[ver[i]] + sze[ver[i]];
if (sze[ver[i]] > sze[son[u]]) son2[u] = son[u], son[u] = ver[i];
else if (sze[ver[i]] > sze[son2[u]]) son2[u] = ver[i];
}
}
inline void getans(int u, int val, int all, int &ret) {
chkmin(ret, val); int v = son[u];
if (v == cut || sze[son[u]] < sze[son2[u]]) v = son2[u];
if (!v) return;
if (2 * sze[v] > all) getans(v, val + all - 2 * sze[v], all, ret);
}
inline void DfsSec(int u) {
for (int i = head[u]; i; i = nxt[i]) {
if (ver[i] == fat[u]) continue;
cut = ver[i]; int x = INT_MAX, y = INT_MAX;
for (int pos = u; pos; pos = fat[pos]) sze[pos] -= sze[ver[i]];
getans(1, g[1] - g[ver[i]] - dep[ver[i]] * sze[ver[i]], sze[1], x);
getans(ver[i], g[ver[i]], sze[ver[i]], y); chkmin(ans, x + y);
for (int pos = u; pos; pos = fat[pos]) sze[pos] += sze[ver[i]];
DfsSec(ver[i]);
}
}
signed main(void) {
// file("");
read(n);
for (int i = 1, u, v; i < n; i++)
read(u), read(v), AddEdge(u, v);
for (int i = 1; i <= n; i++) read(w[i]);
DfsFir(1); DfsSec(1); writeln(ans);
// fwrite(pf, 1, o1 - pf, stdout);
return 0;
}