一次根树遍历求每条边两侧规模,按长度乘规模差累加费用。
OJ: luogu
题目 ID: P2052
难度:普及+/提高-
标签:树形 DP子树大小前向星python
日期: 2026-07-17 02:00
题意
树边费用为边长乘以两侧国家数量之差的绝对值,求总费用。
思路
任选 1 为根。若父子边下方子树大小为 size,两侧规模分别是 size 与 n-size,边贡献就是 weight * abs(n-2*size)。后序累加子树大小即可。
Python 知识
- 百万节点使用
array前向星保存边,避免百万个 Python 小列表。 reversed(order[1:])只做一次自底向上汇总。abs直接表达两侧规模差的绝对值。
代码
python
import sys
from array import array
input = sys.stdin.buffer.readline
n = int(input())
head = array("i", [-1]) * (n + 1)
to = array("i", [0]) * (2 * n - 2)
next_edge = array("i", [0]) * (2 * n - 2)
weight = array("q", [0]) * (2 * n - 2)
edge_count = 0
def add_edge(u, v, length):
global edge_count
to[edge_count] = v
weight[edge_count] = length
next_edge[edge_count] = head[u]
head[u] = edge_count
edge_count += 1
for _ in range(n - 1):
u, v, length = map(int, input().split())
add_edge(u, v, length)
add_edge(v, u, length)
parent = array("i", [0]) * (n + 1)
parent_weight = array("q", [0]) * (n + 1)
order = array("i", [1])
index = 0
while index < n:
node = order[index]
index += 1
edge = head[node]
while edge != -1:
neighbor = to[edge]
if neighbor != parent[node]:
parent[neighbor] = node
parent_weight[neighbor] = weight[edge]
order.append(neighbor)
edge = next_edge[edge]
subtree = array("q", [1]) * (n + 1)
answer = 0
for node in reversed(order[1:]):
subtree[parent[node]] += subtree[node]
for node in range(2, n + 1):
answer += parent_weight[node] * abs(n - 2 * subtree[node])
print(answer)原有 C++ 版本仍保留:
cpp
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 1000000 + 5;
const int MAXM = 2000000 + 5;
static int n;
static int head[MAXN];
static int to[MAXM];
static int nxt[MAXM];
static int weight_arr[MAXM];
static int edge_cnt;
static int parent_arr[MAXN];
static int parent_weight[MAXN];
static int subtree_size[MAXN];
static int order_arr[MAXN];
static int stack_arr[MAXN];
void add_edge(int u, int v, int w) {
++edge_cnt;
to[edge_cnt] = v;
weight_arr[edge_cnt] = w;
nxt[edge_cnt] = head[u];
head[u] = edge_cnt;
}
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
cin >> n;
for (int i = 1; i <= n - 1; ++i) {
int u, v, w;
cin >> u >> v >> w;
add_edge(u, v, w);
add_edge(v, u, w);
}
int top = 0;
int order_cnt = 0;
stack_arr[++top] = 1;
parent_arr[1] = 0;
// 先用迭代 DFS 建出父子关系与访问顺序,避免深递归爆栈。
while (top > 0) {
int u = stack_arr[top--];
order_arr[++order_cnt] = u;
for (int e = head[u]; e != 0; e = nxt[e]) {
int v = to[e];
if (v == parent_arr[u]) {
continue;
}
parent_arr[v] = u;
parent_weight[v] = weight_arr[e];
stack_arr[++top] = v;
}
}
long long ans = 0;
for (int i = 1; i <= n; ++i) {
subtree_size[i] = 1;
}
// 逆序回推子树大小,并统计每条父边的贡献。
for (int i = order_cnt; i >= 2; --i) {
int u = order_arr[i];
long long diff = llabs(1LL * n - 2LL * subtree_size[u]);
ans += 1LL * parent_weight[u] * diff;
subtree_size[parent_arr[u]] += subtree_size[u];
}
cout << ans << '\n';
return 0;
}复杂度
时间 O(n),空间 O(n)。
总结
树边分割问题通常只需要知道一侧子树大小,另一侧就是 n-size。
