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

简单树链剖分 + ST 表题。

因为 $c_i$ 很小，所以考虑状态压缩，用一个二进制数表示一条路径上的 $c_i$ 构成的集合，如果第 $i$ 位为 $1$ 就代表第 $c_i-1$ 种零食包含在该路径上；于是查一条路径上的并集就变成了查一条路径上所有二进制数按位或（$\mathrm{or}$）的 $\mathrm{popcount}$（$\mathrm{popcount}$ 表示一个二进制数中 $1$ 的个数，这里的含义即为该路径上包含的零食种类数）。

树上路径查询考虑用树链剖分维护，而 $\mathrm{or}$ 操作满足消去律，所以静态区间 $\mathrm{or}$ 可以使用 ST 表维护。

参考代码（GNU C++17）：

#include<bits/stdc++.h>
using namespace std;
inline int op(int x,int y){return x|y;}
// 按位或操作
namespace IAOI_lib{
  template<typename T,T(*op)(T,T)> class sparse_table{
    private:
      vector<vector<T> > s;
    public:
      void reset(vector<T> a){
        int k=__lg(a.size());
        s.resize(a.size(),vector<T>(k+1));
        for(int i=0;i<a.size();i++)
          s[i][0]=a[i];
        for(int i=1;i<=k;i++)
          for(int j=0;j+(1<<i)<=a.size();j++)
            s[j][i]=op(s[j][i-1],s[j+(1<<i-1)][i-1]);
      }
      sparse_table(){}
      sparse_table(vector<T> a){reset(a);}
      T query(int l,int r){
        int k=__lg(r-l+1);
        return op(s[l][k],s[r-(1<<k)+1][k]);
      }
  };
} // ST 表模板
class hld{
  private:
    int n;
    vector<int> f,d,e,l,h,t;
    IAOI_lib::sparse_table<int,op> s;
  public:
    hld(vector<vector<int> > &g,vector<int> &w){
      n=w.size(),f.resize(n),d.resize(n),e.resize(n,1);
      h.resize(n,-1),t.resize(n),l.resize(n);
      function<void(int,int)> dfs=[&](int u,int p){
        f[u]=p,d[u]=d[p]+1; int mx=0;
        for(int i:g[u])
          if(i!=p)if(dfs(i,u),e[u]+=e[i];e[i]>mx)mx=e[i],h[u]=i;
      };
      vector<int> a;
      function<void(int,bool)> decomp=[&](int u,bool b){
        t[u]=(b?t[f[u]]:u),l[u]=a.size(),a.emplace_back(w[u]);
        if(~h[u])decomp(h[u],true);
        for(int i:g[u])
          if(i!=f[u]&&i!=h[u])decomp(i,false);
      };
      dfs(0,0),decomp(0,false),s.reset(a);
    }
    inline int query(int u,int v){
      int c=0;
      while(t[u]!=t[v]){
        if(d[t[u]]>d[t[v]])swap(u,v);
        c|=s.query(l[t[v]],l[v]);
        v=f[t[v]];
      }
      if(d[u]>d[v])swap(u,v);
      return c|s.query(l[u],l[v]);
    } // 链查询
}; // 树链剖分模板
int main(){
  ios::sync_with_stdio(false);
  int n,q; cin>>n>>q;
  vector<int> w(n);
  for(auto &i:w)cin>>i,i=1<<i-1;
  // 初始化所有结点的值
  // 第 i 个结点的值为 1<<(c[i]-1)
  vector<vector<int> > g(n);
  for(int i=1;i<n;i++){
    int u,v; cin>>u>>v;
    g[--u].emplace_back(--v);
    g[v].emplace_back(u);
  }
  hld h(g,w);
  while(q--){
    int u,v; cin>>u>>v;
    cout<<__builtin_popcount(h.query(u-1,v-1))<<'\n';
  } // 处理查询
  return 0;
}