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

· · 题解

题目要我们求 s_i 至 t_i 的路径，让我们很容易想到先求出它们的 LCA，然后再去做一些操作。

我的做法：先预处理出一个 cnt 的二维数组，其中 cnt[i][j] 表示从根节点（自己设定）出发，到 i 这个点这条路上 j 这种商品的出现个数。

这个 cnt 用一趟 DFS 就求好了。DFS 到 x 这个点时，枚举它的所有子节点 u。先把它的父节点（x）的 cnt 复制过来，再加上 u 自己的商品。即：

for(int i=1;i<=22;i++){
    cnt[u][i]=cnt[x][i];
}
cnt[u][c[u]]++;

接着就是重点了。我们求 x 到 y 这条路上的答案，可以用到类似树上差分的做法。记 q=\operatorname{lca}_{x,y} 自己手推一下，可以发现答案为：

cnt_{x,i}+cnt_{y,i}-2\times cnt_{\operatorname{t},i}+[c_t=i]

其中 i 是需要枚举的，从 1 到 20。上面公式中的中括号是判断 x 与 y 的 LCA 是否正好有这个商品，如果有就得加 1。这算是一种特殊情况。

不会 LCA 的请移步 P3379。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=1e5+5,K=22;
int n,q,c[N],u,v,s,t,st[N][K+2],dep[N],cnt[N][K+2];
vector<int> g[N];
void dfs(ll x,ll fx){
    dep[x]=dep[fx]+1,st[x][0]=fx;
    for(ll i=1;i<=K;i++){
        st[x][i]=st[st[x][i-1]][i-1];
    }
    for(ll u:g[x]){
        if(u!=fx) dfs(u,x);
    }
}
int LCA(int x,int y){
    if(dep[x]<dep[y]) swap(x,y);
    for(ll i=K;i>=0;i--){
        ll fx=st[x][i];
        if(dep[fx]>=dep[y]) x=fx;
    }
    if(x==y) return x;
    for(int i=20;i>=0;i--){
        ll fa=st[x][i],fb=st[y][i];
        if(fa!=fb) x=fa,y=fb;
    }
    return st[x][0];
}
void DFS(int x,int fx){
    for(int u:g[x]){
        if(u==fx||u==x) continue;
        for(int i=1;i<=22;i++){
            cnt[u][i]=cnt[x][i];
        }
        cnt[u][c[u]]++;
        DFS(u,x);
    }
}
int main(){
    cin>>n>>q;
    for(int i=1;i<=n;i++){
        cin>>c[i];
    }
    for(int i=1;i<n;i++){
        cin>>u>>v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1,1);
    cnt[1][c[1]]++;
    DFS(1,1);
    while(q--){
        cin>>s>>t;
        int ans=0,lc=LCA(s,t);
        for(int i=1;i<=22;i++){
            int tt=0;
            if(c[lc]==i) tt=1;
            if(cnt[s][i]+cnt[t][i]-2*cnt[lc][i]+tt>0) ans++;
        }
        cout<<ans<<'\n';
    }
    return 0;
}