题意:给一个有N个顶点的树,每个顶点上有Ti个元素,树的每条边都有一个长度Ci。假设Cost_i为树中的每个元素到顶点i的距离之和。题目要求找出树中的一个顶点x,使得Cost_x最小。
解法:最近这类题目比较多,解法是通过两次dfs在O(n)的时间复杂度内求出每个点的Cost。具体来说,假设以顶点1为根对树进行dfs。第一次dfs计算出每个子树下面的元素个数Cunt和每个元素到顶点1的距离之和。第二次dfs再计算出每个顶点的Cost。第一次dfs比较简单,关键是第二次dfs。如果已经计算出顶点x的Cost_x,现在要计算和x相连的顶点y的Cost_y。假设一共有Total个元素,x和y相连的边的权值为Length,以顶点y为根的子树所包含的元素个数为Cunt_y。则可以看出:相对于Cost_x,有Total-Cunt_y个元素要多走Length的距离,有Cunt_y个元素要少走Length的距离。所以Cost_y=Cost_x-(Total-2*Cunt_y)*Length。所以可以通过第二遍dfs求出所有的Cost。
注意:由于n比较大,有比较极端的数据,递归深度会非常大。如果用递归写dfs,会RE,所以不能用递归。现在已经越来越多的题卡这一点了...
#include <stdio.h>
#include <memory.h>
const int maxn = 100005;
struct Graph {
int hed[maxn], nxt[maxn*2], pnt[maxn*2], val[maxn*2];
int idx;
void init(int n) {
memset(hed + 1, -1, n * 4);
idx = 0;
}
void addedge(int x, int y, int v) {
pnt[idx] = y; val[idx] = v; nxt[idx] = hed[x]; hed[x] = idx++;
pnt[idx] = x; val[idx] = v; nxt[idx] = hed[y]; hed[y] = idx++;
}
} gra;
struct Node {
int r, x, p;
} sta[maxn];
int teams[maxn], total;
__int64 cunt[maxn], cost[maxn];
void dfs_1() {
int x, y, r, p, top = 0;
sta[0].x = 1; sta[0].r = 0; sta[0].p = gra.hed[1];
cunt[1] = teams[1];
cost[1] = 0;
while (true) {
p = sta[top].p;
if (p == -1) {
top--;
if (top >= 0) {
p = sta[top].p;
x = sta[top].x;
y = gra.pnt[p];
cunt[x] += cunt[y];
cost[x] += cunt[y] * gra.val[p] + cost[y];
sta[top].p = gra.nxt[p];
} else {
break;
}
} else {
x = sta[top].x;
r = sta[top].r;
y = gra.pnt[p];
if (y != r) {
++top;
cunt[y] = teams[y];
cost[y] = 0;
sta[top].r = x;
sta[top].x = y;
sta[top].p = gra.hed[y];
} else {
sta[top].p = gra.nxt[p];
}
}
}
}
void dfs_2() {
int x, y, r, p, top = 0;
sta[0].x = 1; sta[0].r = 0; sta[0].p = gra.hed[1];
while (true) {
p = sta[top].p;
if (p == -1) {
top--;
if (top >= 0) {
p = sta[top].p;
sta[top].p = gra.nxt[p];
} else {
break;
}
} else {
x = sta[top].x;
r = sta[top].r;
y = gra.pnt[p];
if (y != r) {
++top;
cost[y] = cost[x] + (total - 2 * cunt[y]) * gra.val[p];
sta[top].r = x;
sta[top].x = y;
sta[top].p = gra.hed[y];
} else {
sta[top].p = gra.nxt[p];
}
}
}
}
/*
void dfs_1(int x, int r) {
cost[x] = 0;
cunt[x] = teams[x];
for (int p = gra.hed[x]; p != -1; p = gra.nxt[p]) {
int y = gra.pnt[p];
if (y != r) {
dfs_1(y, x);
cunt[x] += cunt[y];
cost[x] += cunt[y] * gra.val[p] + cost[y];
}
}
}
void dfs_2(int x, int r) {
for (int p = gra.hed[x]; p != -1; p = gra.nxt[p]) {
int y = gra.pnt[p];
if (y != r) {
cost[y] = cost[x] + (total - 2 * cunt[y]) * gra.val[p];
dfs_2(y, x);
}
}
}
*/
int main() {
int n, x, y, v;
while (scanf("%d", &n) != EOF) {
total = 0;
for (int i = 1; i <= n; i++) {
scanf("%d", teams + i);
total += teams[i];
}
gra.init(n);
for (int i = 1; i < n; i++) {
scanf("%d %d %d", &x, &y, &v);
gra.addedge(x, y, v);
}
dfs_1();
dfs_2();
__int64 ans = cost[1];
for (int i = 2; i <= n; i++) {
if (cost[i] < ans)
ans = cost[i];
}
printf("%I64d\n", ans);
}
return 0;
}