[NOI2011] 道路修建

GitHub跳转原题关系图返回列表

一次根树遍历求每条边两侧规模,按长度乘规模差累加费用。

OJ: luogu

题目 ID: P2052

难度:普及+/提高-

标签:树形 DP子树大小前向星python

日期: 2026-07-17 02:00

题意

树边费用为边长乘以两侧国家数量之差的绝对值,求总费用。

思路

任选 1 为根。若父子边下方子树大小为 size,两侧规模分别是 sizen-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