AcWing 252. 树 题解

题目分析

本题给定一棵有(N)个节点的树，每条边都有权值，要求计算树上长度不超过(K)的路径数量。输入包含多组测试用例，当输入(N = 0)，(K = 0)时结束输入。

解题思路

本题采用树的重心分治算法来解决。树的重心分治可以有效地将树分解，减少计算量，避免时间复杂度退化为(O(n^2))。
1. 构建邻接表：通过add函数构建树的邻接表，方便后续对树进行遍历和操作。
2. 计算子树大小：使用get_size函数计算以某个节点为根的子树大小，用于后续找树的重心。
3. 寻找树的重心：利用get_wc函数找到树的重心，选择重心作为分治的中心点可以保证每次分治后子树的大小尽量平衡，从而优化时间复杂度。
4. 计算路径长度：get_dist函数用于计算从某个节点到其所有子节点的路径长度，并存储在数组中。
5. 统计满足条件的路径数量：get函数对路径长度数组进行排序，然后统计满足长度不超过(K)的路径数量。
6. 分治计算：calc函数以树的重心为分界点，递归地计算满足条件的路径数量，同时去除重复计算的部分。

代码逐段分析

1. 头文件和全局变量定义

#include <iostream>
#include <cstring>
#include <algorithm>

using namespace std;

const int N = 10010, M = N * 2;

int n, m;
int h[N], e[M], w[M], ne[M], idx;
bool st[N];
int p[N], q[N];

• 引入必要的头文件，包括输入输出、字符串操作和算法相关的头文件。
• 定义常量(N)表示树的节点数，(M)表示边数的两倍（因为是无向图）。
• 变量(n)表示节点数，(m)表示路径长度的上限(K)。
• h[N]是邻接表的头指针数组，e[M]存储边的另一端节点，w[M]存储边的权值，ne[M]指向下一条边的下标，idx用于记录边的编号。
• st[N]用于标记节点是否已被处理（在分治过程中删除）。

• p[N]和q[N]用于存储路径长度。

• 添加边函数

void add(int a, int b, int c)
{
    e[idx] = b, w[idx] = c, ne[idx] = h[a], h[a] = idx ++ ;
}

• 该函数用于向邻接表中添加一条边，从节点a到节点b，权值为c。

• 计算子树大小函数

int get_size(int u, int fa)  // 求子树大小
{
    if (st[u]) return 0;
    int res = 1;
    for (int i = h[u]; ~i; i = ne[i])
        if (e[i] != fa)
            res += get_size(e[i], u);
    return res;
}

• 该函数递归计算以节点u为根的子树大小。如果节点u已被处理（st[u]为true），则返回(0)。否则，初始子树大小为(1)（包含节点u自身），然后遍历节点u的所有子节点（排除父节点fa），递归计算子树大小并累加。

• 寻找树的重心函数

int get_wc(int u, int fa, int tot, int& wc)  // 求重心
{
    if (st[u]) return 0;
    int sum = 1, ms = 0;
    for (int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        if (j == fa) continue;
        int t = get_wc(j, u, tot, wc);
        ms = max(ms, t);
        sum += t;
    }
    ms = max(ms, tot - sum);
    if (ms <= tot / 2) wc = u;
    return sum;
}

• 该函数用于寻找以节点u为根的子树的重心。如果节点u已被处理，则返回(0)。
• 初始化子树大小sum为(1)，最大子树大小ms为(0)。遍历节点u的所有子节点（排除父节点fa），递归计算子树大小t，更新ms为已遍历子树中的最大子树大小，并累加子树大小sum。

• 计算以节点u为根的子树中，除当前子树外的其他部分的大小（tot - sum），并更新ms。如果ms不超过子树总大小tot的一半，则节点u为重心，将其赋值给wc。最后返回子树大小sum。

• 计算路径长度函数

void get_dist(int u, int fa, int dist, int& qt)
{
    if (st[u]) return;
    q[qt ++ ] = dist;
    for (int i = h[u]; ~i; i = ne[i])
        if (e[i] != fa)
            get_dist(e[i], u, dist + w[i], qt);
}

• 该函数用于计算从节点u到其所有子节点的路径长度，并将路径长度存储在数组q中。如果节点u已被处理，则返回。

• 将当前节点u到根节点的路径长度dist存入数组q，然后遍历节点u的所有子节点（排除父节点fa），递归计算子节点到根节点的路径长度（dist + w[i]）。

• 统计满足条件的路径数量函数

int get(int a[], int k)
{
    sort(a, a + k);
    int res = 0;
    for (int i = k - 1, j = -1; i >= 0; i -- )
    {
        while (j + 1 < i && a[j + 1] + a[i] <= m) j ++ ;
        j = min(j, i - 1);
        res += j + 1;
    }
    return res;
}

• 该函数对路径长度数组a进行排序，然后统计满足路径长度之和不超过m的路径数量。

• 从数组末尾开始遍历，使用双指针法，通过移动指针j找到满足条件的路径对数，并累加到结果res中。

• 分治计算函数

int calc(int u)
{
    if (st[u]) return 0;
    int res = 0;
    get_wc(u, -1, get_size(u, -1), u);
    st[u] = true;  // 删除重心

    int pt = 0;
    for (int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i], qt = 0;
        get_dist(j, -1, w[i], qt);
        res -= get(q, qt);
        for (int k = 0; k < qt; k ++ )
        {
            if (q[k] <= m) res ++ ;
            p[pt ++ ] = q[k];
        }
    }
    res += get(p, pt);

    for (int i = h[u]; ~i; i = ne[i]) res += calc(e[i]);
    return res;
}

• 该函数以节点u为起点进行分治计算。如果节点u已被处理，则返回(0)。
• 首先找到节点u所在子树的重心，并标记重心已被处理（从树中删除）。
• 遍历重心的所有子节点，计算每个子节点到其所有子节点的路径长度，统计这些路径中不满足条件的数量（减去重复计算的部分），并将路径长度存储在数组p中。
• 统计数组p中满足条件的路径数量，并累加到结果res中。

• 递归地对重心的所有子节点进行分治计算，并将结果累加到res中。最后返回满足条件的路径总数res。

• 主函数

int main()
{
    while (scanf("%d%d", &n, &m), n || m)
    {
        memset(st, 0, sizeof st);
        memset(h, -1, sizeof h);
        idx = 0;
        for (int i = 0; i < n - 1; i ++ )
        {
            int a, b, c;
            scanf("%d%d%d", &a, &b, &c);
            add(a, b, c), add(b, a, c);
        }

        printf("%d\n", calc(0));
    }

    return 0;
}

• 主函数通过循环读取输入，当输入的n和m都为(0)时结束循环。
• 每次读取一组测试用例后，初始化标记数组st和邻接表头指针数组h，并构建树的邻接表。
• 调用calc函数计算满足条件的路径数量，并输出结果。

时间复杂度分析

• 每次寻找重心和计算子树大小的时间复杂度为(O(n))，由于树的重心分治每次能将树分成两部分，递归次数为(O(\log n))，所以这部分总时间复杂度为(O(n \log n))。
• 统计路径数量时，排序的时间复杂度为(O(n \log n))，双指针统计的时间复杂度为(O(n))，总体统计路径数量的时间复杂度也为(O(n \log n))。
• 因此，本题的时间复杂度为(O(n \log n))。

空间复杂度分析

主要空间消耗在于邻接表相关数组、标记数组和存储路径长度的数组等，空间复杂度为(O(n))。