AcWing 2226. 开店 题解

题目分析

本题设定在一个树形结构的幻想乡地图中，每个地方住着一个妖怪，妖怪有年龄属性。需要在某个地方(u)开一家面向年龄在(L)到(R)之间妖怪的店，并求出这些妖怪到点(u)的距离之和（方便值）。树的顶点度数不超过(3)，给定多个开店方案，要求计算每个方案的方便值。

解题思路

本题通过树的重心分治算法来优化计算过程，同时借助一些辅助数据结构记录路径和妖怪年龄信息，从而高效地计算出每个开店方案的方便值。
1. 构建邻接表：利用add函数构建树的邻接表，用于后续对树进行遍历和操作。
2. 计算子树大小：get_size函数递归计算以某个节点为根的子树大小，为寻找树的重心提供依据。
3. 寻找树的重心：get_wc函数找到树或子树的重心，使每次分治时子树大小尽量平衡，降低时间复杂度。
4. 计算路径和妖怪信息：get_dist函数从某个节点出发，计算到其所有子节点的路径距离，并记录该路径所属的重心以及相关妖怪的年龄和距离信息。
5. 分治预处理：calc函数以树的重心为分界点，递归处理每个子树。在处理过程中，整理每个节点到其不同方向子节点的妖怪年龄和距离信息，并进行排序和前缀和处理，方便后续查询。
6. 查询计算：query函数根据给定的开店地点(u)和年龄范围(l)到(r)，利用之前记录的信息，计算出满足条件的妖怪到(u)的距离之和。

代码逐段分析

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

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

using namespace std;

typedef long long LL;

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

int n, m, A;
int h[N], e[M], w[M], ne[M], idx;
int age[N];
bool st[N];
struct Father
{
    int u, num;
    LL dist;
};
vector<Father> f[N];
struct Son
{
    int age;
    LL dist;
    bool operator< (const Son& t) const
    {
        return age < t.age;
    }
};
vector<Son> son[N][3];

• 引入必要的头文件，包括输入输出、字符串操作和算法相关的头文件。
• 定义LL为long long的别名。
• 定义常量(N)表示树的节点数，(M)表示边数的两倍（因为是无向图）。
• 变量(n)表示节点数，(m)表示开店方案的数量，(A)是一个用于计算年龄范围的参数。
• h[N]是邻接表的头指针数组，e[M]存储边的另一端节点，w[M]存储边的权值，ne[M]指向下一条边的下标，idx用于记录边的编号。
• age[N]存储每个节点处妖怪的年龄。
• st[N]用于标记节点是否已被处理（在分治过程中删除）。
• 定义结构体Father，用于记录节点到重心的信息，包括节点编号u、所在子树编号num和到重心的距离dist。f[N]是一个向量数组，存储每个节点到不同重心的信息。

• 定义结构体Son，用于记录子节点的妖怪年龄和到当前节点的距离，son[N][3]是一个二维向量数组，存储每个节点不同方向子节点的相关信息。

• 添加边函数

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, LL dist, int wc, int k, vector<Son>& p)
{
    if (st[u]) return;
    f[u].push_back({wc, k, dist});
    p.push_back({age[u], dist});
    for (int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        if (j == fa) continue;
        get_dist(j, u, dist + w[i], wc, k, p);
    }
}

• 该函数从节点u出发，计算到其所有子节点的路径距离dist。如果节点u已被处理，则返回。
• 将节点u到重心wc的信息（重心编号、子树编号、距离）存入f[u]，将节点u处妖怪的年龄和到当前节点的距离存入p。

• 然后遍历节点u的所有子节点（排除父节点fa），递归计算子节点的路径和妖怪信息。

• 分治预处理函数

void calc(int u)
{
    if (st[u]) return;
    get_wc(u, -1, get_size(u, -1), u);
    st[u] = true;

    for (int i = h[u], k = 0; ~i; i = ne[i])
    {
        int j = e[i];
        if (st[j]) continue;
        auto& p = son[u][k];
        p.push_back({-1, 0}), p.push_back({A + 1, 0});
        get_dist(j, -1, w[i], u, k, p);
        k ++ ;
        sort(p.begin(), p.end());
        for (int i = 1; i < p.size(); i ++ ) p[i].dist += p[i - 1].dist;
    }

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

• 该函数以节点u为起点进行分治计算。如果节点u已被处理，则返回。
• 首先找到节点u所在子树的重心，并标记重心已被处理（从树中删除）。
• 遍历重心的所有子节点：
  • 初始化子节点信息向量p，并添加两个哨兵元素（年龄为(-1)和(A + 1)，距离为(0)）。
  • 调用get_dist函数获取子节点的妖怪年龄和距离信息。
  • 对获取到的信息按照妖怪年龄进行排序，并计算距离的前缀和，方便后续查询。

• 递归地对重心的所有子节点进行分治计算。

• 查询计算函数

LL query(int u, int l, int r)
{
    LL res = 0;
    for (auto& t: f[u])
    {
        int g = age[t.u];
        if (g >= l && g <= r) res += t.dist;
        for (int i = 0; i < 3; i ++ )
        {
            if (i == t.num) continue;
            auto& p = son[t.u][i];
            if (p.empty()) continue;
            int a = lower_bound(p.begin(), p.end(), Son({l, -1})) - p.begin();
            int b = lower_bound(p.begin(), p.end(), Son({r + 1, -1})) - p.begin();
            res += t.dist * (b - a) + p[b - 1].dist - p[a - 1].dist;
        }
    }

    for (int i = 0; i < 3; i ++ )
    {
        auto& p = son[u][i];
        if (p.empty()) continue;
        int a = lower_bound(p.begin(), p.end(), Son({l, -1})) - p.begin();
        int b = lower_bound(p.begin(), p.end(), Son({r + 1, -1})) - p.begin();
        res += p[b - 1].dist - p[a - 1].dist;
    }

    return res;
}

• 该函数根据给定的开店地点u和年龄范围l到r，计算满足条件的妖怪到u的距离之和。
• 遍历节点u到不同重心的信息：
  • 如果当前节点的妖怪年龄在范围内，将其到重心的距离累加到结果res中。
  • 对于当前节点不同方向的子节点信息（排除当前子树方向），通过二分查找找到年龄范围对应的起始和结束位置a和b，计算这些子节点中满足年龄范围的妖怪到u的距离并累加到res中。
• 再遍历节点u自身不同方向的子节点信息，同样通过二分查找计算满足年龄范围的妖怪到u的距离并累加到res中。

• 最后返回计算得到的方便值res。

• 主函数

int main()
{
    scanf("%d%d%d", &n, &m, &A);
    for (int i = 1; i <= n; i ++ ) scanf("%d", &age[i]);
    memset(h, -1, sizeof h);
    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);
    }
    calc(1);
    LL res = 0;
    while (m -- )
    {
        int u, a, b;
        scanf("%d%d%d", &u, &a, &b);
        int l = (a + res) % A, r = (b + res) % A;
        if (l > r) swap(l, r);
        res = query(u, l, r);
        printf("%lld\n", res);
    }

    return 0;
}

• 主函数读入树的节点数n、开店方案数量m和参数A，以及每个节点处妖怪的年龄。
• 初始化邻接表头指针数组h，并根据输入构建树的邻接表。
• 从节点(1)开始调用calc函数进行分治预处理。
• 对于每个开店方案：
  • 读入开店地点u和年龄范围参数a、b。
  • 计算实际的年龄范围l和r（通过与上一个方案的结果res进行计算并取模），如果l > r则交换。
  • 调用query函数计算当前方案的方便值res并输出。

时间复杂度分析

• 每次寻找重心和计算子树大小的时间复杂度为(O(n))，由于树的重心分治每次能将树分成两部分，递归次数为(O(\log n))，所以这部分总时间复杂度为(O(n \log n))。
• 在预处理和查询过程中，对于每个节点，处理其子节点信息的时间复杂度与子节点数量和二分查找有关，总体时间复杂度也为(O(n \log n))。
• 因此，本题的时间复杂度为(O(n \log n + m \log n))，其中(m)是开店方案的数量。

空间复杂度分析

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