[蓝桥杯 2024 省 A] 零食采购 题解

题目描述

小蓝准备去星际旅行，出发前想在本星系采购一些零食，星系内有 $n$ 颗星球，由 $n-1$ 条航路连接为连通图，第 $i$ 颗星球卖第 $c_i$ 种零食特产。小蓝想出了 $q$ 个采购方案，第 $i$ 个方案的起点为星球 $s_i$ ，终点为星球 $t_i$ ，对于每种采购方案，小蓝将从起点走最短的航路到终点，并且可以购买所有经过的星球上的零食（包括起点终点），请计算每种采购方案最多能买多少种不同的零食。

输入格式

输入的第一行包含两个正整数 $n$，$q$，用一个空格分隔。
第二行包含 $n$ 个整数 $c_1,c_2,\cdots ,c_n$，相邻整数之间使用一个空格分隔。
接下来 $n-1$ 行，第 $i$ 行包含两个整数 $u_i,v_i$，用一个空格分隔，表示一条
航路将星球 $u_i$ 与 $v_i$ 相连。
接下来 $q$ 行，第 $i$ 行包含两个整数 $s_i
, t_i $，用一个空格分隔，表示一个采购方案。

输出格式

输出 $q$ 行，每行包含一个整数，依次表示每个采购方案的答案。

输入输出样例 #1

输入 #1

4 2
1 2 3 1
1 2
1 3
2 4
4 3
1 4

输出 #1

3
2

说明/提示

第一个方案路线为 {4,2,1,3}\{4, 2, 1, 3\}{4,2,1,3}，可以买到第 $1,2,3$ 种零食；
第二个方案路线为 {1,2,4}\{1, 2, 4\}{1,2,4}，可以买到第 $1,2$ 种零食。

对于 20% 的评测用例，$1 ≤ n, q ≤ 5000 $；
对于所有评测用例，$1\le n,q\le 10^5，1\le c_i\le 20，1\le u_i,v_i\le n，1\le s_i,t_i\le n$。

预备知识

在解这道题之前，我们需要有一定的求最近公共祖先（Least Common Ancestors，LCA）的知识。
最近公共祖先简称 LCA。两个节点的最近公共祖先，就是这两个点的公共祖先里面，离根最远的那个。

如图，4和5的最近公共祖先就是2，8和4的最近公共祖先就是2。
求LCA的算法有很多，个人认为比较好的算法是倍增算法。

倍增算法

何为倍增？
所谓倍增，就是按$2$的倍数来增大，也就是跳$1,2,4,8,\mathrm{16...}$，不过在这我们不是按从小到大跳，而是从大向小跳。如果大的跳不过去，再把它调小。这是因为从小开始跳，可能会出现“悔棋”的现象。用$5$为例，从小向大跳，$5\ne 1+2+4$,所以我们还要回溯一步，然后才能得出$5\ne 1+4$；而从大向小跳，直接可以得出$5\ne 1+4$。这也可以拿二进制为例，$5(101)$，从高位向低位填很简单，如果填了这位之后比原数大了，那就不填。
这个算法使向上跳的次数大大减小。这个算法的时间复杂度为$O(nlogn)$。
算法实现
想要实现这个算法，首先要记录各个点的深度和该点向上跳$2^i$层的祖先，用数组$dep[x]$表示x的深度，$dp[x][i]$表示结点向上跳$2^i$层的祖先，具体的代码实现如下：

// DFS预处理每个节点的深度、倍增祖先数组
// x: 当前节点，f: 父节点
void dfs(int x, int f) {
    dep[x] = dep[f] + 1;  // 计算当前节点深度（父节点深度+1）
    // 初始化倍增祖先数组（dp[x][i]表示x的2^i级祖先）
    dp[x][0] = f;  // 直接父节点（2^0级祖先）
    // 预处理倍增数组（二进制跳跃法）
    for (int i = 1; i < 20; ++i) {
        dp[x][i] = dp[dp[x][i - 1]][i - 1];  // 2^i = 2^(i-1) + 2^(i-1)
        if (dp[x][i] == 0) break;  // 如果祖先不存在则停止预处理
    }
    // 递归处理子节点
    for (auto ch : child[x]) {
        if (ch == f) continue;  // 跳过父节点
        dfs(ch, x);  // 深度优先遍历
    }
}

接下来就是倍增LCA了，我们先把两个点提到同一高度，再统一开始跳。使用逆序循实现二进制位的从高到低检查（跳跃）。具体的代码实现和一些注意事项如下：

// 倍增法求解最近公共祖先（LCA）
int LCA(int x, int y) {
    // 确保x是较深的节点
    if (dep[x] < dep[y]) swap(x, y);
    
    // 将x上提到与y同深度（二进制跳跃法）
    for (int i = 19; i >= 0; --i) {
        if (dep[dp[x][i]] >= dep[y]) {  // 如果跳跃后仍深于y
            x = dp[x][i];             // 执行跳跃
        }
    }
    if (x == y) return x;  // 如果此时相同，则说明y一定是x的祖先，即x和y的LCA为y
    // 同时上提x和y直到找到共同祖先
    for (int i = 19; i >= 0; --i) {
    //这里需要一定的理解力，如果他们跳之后相同，则说明是他们的公共祖先，但不一定是最近公共祖先。
    //比如6和10，在跳的时候，会直接跳到1，但1只是它们的祖先，它们的LCA其实是3。所以我们要跳到它们LCA的下面一层，
        if (dp[x][i] != dp[y][i]) {  // 当祖先不同时继续跳跃
            x = dp[x][i];
            y = dp[y][i];
        }
    }
    
    return dp[x][0];  // 最终返回公共祖先的父节点（即LCA）
}
在这里插入代码片

倍增算法就是这样，还是很好理解的。

解题思路

回到本题，小蓝将从起点走最短的航路到终点。因为题目中说了“星系内有 $n$ 颗星球，由 $n-1$ 条航路连接为连通图”，所以这$n$个点其实构成了一棵树，每两个星球之间仅存在一条最短的航路，而且一定是经过它们的最近公共祖先，这里就用到了前面的预备知识。其实到这一步，这道题的答案基本就出来了。
我们可以从根节点（定为1）出发，购买零食。对于从1出发，到每个星球能够购买的零食种类和数量可以用一个数组$nums[x][i],i\in [0,20]$记录。从星球$u$和星球$v$到它们的最近公共祖先$f$，分别能够买到的零食种类和数量可以表示为$nums[u][i]-nums[dp[f][0]][i],i\in [0,20]$，$nums[v][i]-nums[dp[f][0]][i],i\in [0,20]$。然后减去$nums[f][i],i\in [0,20]$，最后统计不为0的数量即为最后结果。
完整代码如下：

#include <iostream>
#include <vector>
using namespace std;

int nums[100001][20]; // nums[x][c] 表示从根到节点x的路径上颜色c出现的次数
int c[100001];        // 存储每个节点的颜色值（0~19）
int dep[100001];      // 存储每个节点的深度
int dp[100001][20];   // 倍增祖先数组，dp[x][i]表示x的2^i级祖先
int n, q, u, v, s, t, ans;
vector<int> child[100001]; // 树的邻接表存储

// DFS预处理每个节点的深度、倍增祖先数组
// x: 当前节点，f: 父节点
void dfs(int x, int f) {
    // 继承父节点的颜色统计（建立前缀和）
    for (int i = 0; i < 20; ++i)
        nums[x][i] = nums[f][i];
    nums[x][c[x]]++; // 当前节点颜色计数+1
    dep[x] = dep[f] + 1;  // 计算当前节点深度（父节点深度+1）
    // 初始化倍增祖先数组（dp[x][i]表示x的2^i级祖先）
    dp[x][0] = f;  // 直接父节点（2^0级祖先）
    // 预处理倍增数组（二进制跳跃法）
    for (int i = 1; i < 20; ++i) {
        dp[x][i] = dp[dp[x][i - 1]][i - 1];  // 2^i = 2^(i-1) + 2^(i-1)
        if (dp[x][i] == 0) break;  // 如果祖先不存在则停止预处理
    }
    // 递归处理子节点
    for (auto ch : child[x]) {
        if (ch == f) continue;  // 跳过父节点
        dfs(ch, x);  // 深度优先遍历
    }
}

// 倍增法求解最近公共祖先（LCA）
int LCA(int x, int y) {
    // 确保x是较深的节点
    if (dep[x] < dep[y]) swap(x, y);

    // 将x上提到与y同深度（二进制跳跃法）
    for (int i = 19; i >= 0; --i) {
        if (dep[dp[x][i]] >= dep[y]) {  // 如果跳跃后仍深于y
            x = dp[x][i];             // 执行跳跃
        }
    }
    if (x == y) return x;  // 如果此时相同，则说明y一定是x的祖先，即x和y的LCA为y
    // 同时上提x和y直到找到共同祖先
    for (int i = 19; i >= 0; --i) {
        //这里需要一定的理解力，如果他们跳之后相同，则说明是他们的公共祖先，但不一定是最近公共祖先。
        //比如6和10，在跳的时候，会直接跳到1，但1只是它们的祖先，它们的LCA其实是3。所以我们要跳到它们LCA的下面一层，
        if (dp[x][i] != dp[y][i]) {  // 当祖先不同时继续跳跃
            x = dp[x][i];
            y = dp[y][i];
        }
    }

    return dp[x][0];  // 最终返回公共祖先的父节点（即LCA）
}

int main() {
    ios::sync_with_stdio(false); // 关闭同步加速输入输出
    cin.tie(0); cout.tie(0);     // 解除cin/cout绑定
    // 输入处理
    cin >> n >> q;
    for (int i = 1; i <= n; ++i) {
        cin >> c[i];
        c[i]--; // 颜色转为0~19范围
    }
    // 建树
    for (int i = 1; i < n; ++i) {
        cin >> u >> v;
        child[u].push_back(v);
        child[v].push_back(u);
    }
    // 预处理深度、颜色前缀和、倍增数组
    dfs(1, 0); // 以1为根节点初始化
    // 处理查询
    while (q--) {
        cin >> s >> t;
        int lca = LCA(s, t);
        ans = 0;
        // 计算路径上的颜色种类数
        for (int i = 0; i < 20; ++i) {
            // 路径颜色数 = s到根 + t到根 - lca到根 - lca父到根
            int cnt = nums[s][i] + nums[t][i]
                - nums[lca][i] - nums[dp[lca][0]][i];
                if (cnt > 0) ans++; // 该颜色存在时计数
        }
        cout << ans << "\n";
    }
    return 0;
}