[USACO10MAR] Great Cow Gathering G

形式化题目

给定一棵包含 $N$ 个结点的带权树，结点编号为 $1 \sim N$ 。每条边 $i$ 连接农场 $A_i$ 和 $B_i$ ，长度为 $L_i$ 。每个结点 $i$ 上居住着 $C_i$ 头奶牛。

选择任意一个结点 $X$ 作为集会地点，其不方便程度定义为所有奶牛到结点 $X$ 的路程之和：

\mathrm{Cost}(X) = \sum_{i=1}^N C_i \times \mathrm{dist}(i, X)

其中 $\mathrm{dist}(i, X)$ 表示树上结点 $i$ 到结点 $X$ 的简单路径权值和。要求找出最方便的地点，使总不方便程度最小，并输出该最小值。

暴力解法

思路

最直接的做法是枚举每个结点 $X \in [1, N]$ 作为候选集会点。对于确定的候选点 $X$ ：

以 $X$ 为起点进行一次广度优先搜索（BFS）或深度优先搜索（DFS），计算整棵树上每个结点 $i$ 到 $X$ 的距离 $\mathrm{dist}(i, X)$ 。
累加计算当前集会点的总代价 $\mathrm{Cost}(X) = \sum_{i=1}^N C_i \times \mathrm{dist}(i, X)$ 。
对所有 $1 \leqslant X \leqslant N$ 取最小值。

代码

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62#include <bits/stdc++.h>
using namespace std;

const int MAXN = 30;

struct Edge {
    int to;
    int w;
};

int n;
long long cows[MAXN];
vector<Edge> g[MAXN];

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    // brute.cpp：对每个点都做一次最短路/BFS 树上距离统计，直接求总代价。
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> cows[i];
        g[i].clear();
    }

    for (int i = 1; i < n; i++) {
        int u, v, w;
        cin >> u >> v >> w;
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }

    long long ans = (1LL << 62);
    for (int start = 1; start <= n; start++) {
        vector<long long> dist(n + 1, -1);
        queue<int> q;
        q.push(start);
        dist[start] = 0;

        while (!q.empty()) {
            int u = q.front();
            q.pop();
            for (size_t i = 0; i < g[u].size(); i++) {
                int v = g[u][i].to;
                if (dist[v] != -1) {
                    continue;
                }
                dist[v] = dist[u] + g[u][i].w;
                q.push(v);
            }
        }

        long long cost = 0;
        for (int i = 1; i <= n; i++) {
            cost += cows[i] * dist[i];
        }
        ans = min(ans, cost);
    }

    cout << ans << '\n';
    return 0;
}

复杂度

时间复杂度：对每个候选集会点做一次树的遍历需要 $O(N)$ ，共枚举 $N$ 个点，总时间复杂度为 $O(N^2)$ 。
空间复杂度：存储邻接表和距离数组需要 $O(N)$ 空间。

瓶颈

在 $N \leqslant 10^5$ 的约束下， $O(N^2)$ 计算量达到 $10^{10}$ 级别，必然超时。

瓶颈在于：相邻两个结点作为集会点时，绝大多数奶牛走过的路径具有极大的重叠。如果每次都从头做一次全树遍历，完全忽略了相邻点之间代价的递推关联。

正解

思路

这是一道经典的树上换根动态规划（Re-rooting DP）问题。

我们可以先任选一个点（通常固定为结点 $1$ ）作为根结点，将无根树转化为有根树。

1. 第一次遍历：自底向上统计子树贡献

以结点 $1$ 为根进行后序遍历（或自底向上的拓扑顺序），计算两个数组：

$\mathrm{sub}[u]$ ：以 $u$ 为根的子树内的奶牛总数。
$\mathrm{sub}[u] = C_u + \sum_{v \in \mathrm{children}(u)} \mathrm{sub}[v]$
$\mathrm{dist}[1]$ ：若将集会地点选在结点 $1$ 时的全树总代价。在自底向上合并时，子树 $v$ 内的所有奶牛跨过边 $(u, v)$ 前往 $u$ ，每头牛都要多走边权 $w(u, v)$ 的距离，因此：
$\mathrm{dist}[u] = \sum_{v \in \mathrm{children}(u)} \Big(\mathrm{dist}[v] + \mathrm{sub}[v] \times w(u, v)\Big)$
全树总奶牛数即为 $\mathrm{total} = \mathrm{sub}[1]$ 。

2. 第二次遍历：自顶向下换根转移

假设当前已经求出以结点 $u$ 为集会地点时的总代价 $\mathrm{dist}[u]$ 。现在考虑将集会地点沿边 $(u, v)$ 从父结点 $u$ 移动到子结点 $v$ ，边长为 $w$ 。

此时整棵树上的所有奶牛可以分为互斥的两部分：

位于 $v$ 的子树内部的奶牛：共 $\mathrm{sub}[v]$ 头。集会点移到 $v$ 后，它们前往集会点的路程都整体减少了 $w$ ，总代价减少 $\mathrm{sub}[v] \times w$ 。
位于 $v$ 的子树外部的其余奶牛：共 $\mathrm{total} - \mathrm{sub}[v]$ 头。它们要到达新集会点 $v$ ，必须先走到 $u$ 再跨过边 $(u, v)$ 到达 $v$ ，路程都整体增加了 $w$ ，总代价增加 $(\mathrm{total} - \mathrm{sub}[v]) \times w$ 。

两部分变化量合并整理，即可得到换根转移公式：

\begin{aligned} \mathrm{dist}[v] &= \mathrm{dist}[u] - \mathrm{sub}[v] \times w + (\mathrm{total} - \mathrm{sub}[v]) \times w \\ &= \mathrm{dist}[u] + (\mathrm{total} - 2 \times \mathrm{sub}[v]) \times w \end{aligned}

通过自顶向下的前序遍历，利用父结点的答案即可在 $O(1)$ 时间递推出子结点的答案。

样例树以 1 为初始根时, 各节点的子树牛数与集会总代价如下. sub[u] 是以 u 为根的子树牛数; dist[u] 是集会点放在 u 时全体牛的路程和. 第二遍换根从 dist[1]=17 出发: 例如沿 1→3 的边权 1 换根, 3 子树有 3 头牛, 子树内代价减 3、外面 1 头牛代价加 1, 所以 dist[3]=17-3+1=15.

节点 $u$	牛数 $C_u$	子树牛数 `sub[u]`	集会代价 `dist[u]`
1	1	4	17
2	1	1	19
3	0	3	15
4	0	2	15
5	2	2	15

下图聚焦换根时跨过边 $(u,v)$ 的代价变化. v 子树内的牛有 sub[v] 头, 各少走 $w$ ; 子树外的牛有 total-sub[v] 头, 各多走 $w$ .

flowchart LR
 U["旧集会点 u"] -->|"边长 w"| V["新集会点 v"]
 V --> I["v 子树内: sub[v] 头牛<br/>每头少走 w<br/>总代价 -sub[v]·w"]
 U --> O["v 子树外: total-sub[v] 头牛<br/>每头多走 w<br/>总代价 +(total-sub[v])·w"]
 I --> S["合并变化量"]
 O --> S
 S --> F["dist[v]=dist[u]+(total-2·sub[v])·w"]

两部分恰好覆盖全体牛, 所以总变化是 -sub[v]·w + (total-sub[v])·w. 这就是换根公式中系数 total-2·sub[v] 的来源.

两类奶牛在换根时的距离变化关系如下所示：

graph G {
  node [shape=box, style="rounded"];
  U [label="父结点 u"];
  V [label="子结点 v"];
  X [label="v 子树内部的牛\n(路程减少 w)"];
  Y [label="v 子树外部的牛\n(路程增加 w)"];
  U -- V [label="边权 w"];
  V -- X;
  U -- Y;
}

遍历完成后，取所有 $\mathrm{dist}[i]$ （ $1 \leqslant i \leqslant N$ ）的最小值即为答案。

代码

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98#include <bits/stdc++.h>
using namespace std;

const int MAXN = 100000 + 5;

struct Edge {
    int to;
    int w;
};

int n;
long long cows[MAXN];
vector<Edge> g[MAXN];
int parent_arr[MAXN];
int parent_w[MAXN];
long long sub_cows[MAXN];
long long dist_sum[MAXN];

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> cows[i];
        g[i].clear();
        parent_arr[i] = 0;
        parent_w[i] = 0;
        sub_cows[i] = 0;
        dist_sum[i] = 0;
    }

    for (int i = 1; i < n; i++) {
        int u, v, w;
        cin >> u >> v >> w;
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }

    vector<int> order;
    order.reserve(n);
    stack<int> st;
    st.push(1);
    parent_arr[1] = 0;

    while (!st.empty()) {
        int u = st.top();
        st.pop();
        order.push_back(u);
        for (size_t i = 0; i < g[u].size(); i++) {
            int v = g[u][i].to;
            if (v == parent_arr[u]) {
                continue;
            }
            parent_arr[v] = u;
            parent_w[v] = g[u][i].w;
            st.push(v);
        }
    }

    long long total_cows = 0;
    for (int idx = (int)order.size() - 1; idx >= 0; idx--) {
        int u = order[idx];
        sub_cows[u] = cows[u];
        for (size_t i = 0; i < g[u].size(); i++) {
            int v = g[u][i].to;
            int w = g[u][i].w;
            if (v == parent_arr[u]) {
                continue;
            }
            sub_cows[u] += sub_cows[v];
            dist_sum[u] += dist_sum[v] + sub_cows[v] * w;
        }
    }
    total_cows = sub_cows[1];

    // 换根：从 u 换到儿子 v，v 子树里的牛会整体更近 w，
    // 其余牛会整体更远 w。
    for (size_t idx = 0; idx < order.size(); idx++) {
        int u = order[idx];
        for (size_t i = 0; i < g[u].size(); i++) {
            int v = g[u][i].to;
            int w = g[u][i].w;
            if (v == parent_arr[u]) {
                continue;
            }
            dist_sum[v] = dist_sum[u] + (total_cows - 2LL * sub_cows[v]) * w;
        }
    }

    long long ans = dist_sum[1];
    for (int i = 2; i <= n; i++) {
        ans = min(ans, dist_sum[i]);
    }

    cout << ans << '\n';
    return 0;
}

复杂度

时间复杂度：非递归遍历求拓扑序、自底向上后序 DP 和自顶向下换根转移均遍历树上的每条边常数次，总时间复杂度为 $O(N)$ ，可在规定时间内轻松通过。
空间复杂度：邻接表、遍历顺序栈以及各个状态数组均占用线性空间，总空间复杂度为 $O(N)$ 。

总结

解决树上最佳聚会点/带权重心问题的核心是换根思想：

先任意定根，通过一次后序遍历求出基准根结点的答案及各子树权值和；
观察集会点移动一步时，整棵树权值的划分——跨过树边后，子树内的点距离全部 $-w$ ，子树外的点距离全部 $+w$ ；
将 $O(N^2)$ 的全树重新遍历压缩为 $O(1)$ 的状态递推，从而在 $O(N)$ 时间内求出所有顶点的聚会代价。

图示解析

下图完整展示了本题的题目模型、换根推导推演、关键实现与复盘思考：