[蓝桥杯 2023 省 A] 颜色平衡树 题解

· · 题解

本题可以使用莫队思想。

对于树上每个结点，按照 dfs 序打上时间戳。然后每一个点的子树里的答案就可以转化为一个区间里的答案，可以使用莫队。

莫队中的 add 和 del 函数可以使用多重集合维护，集合中的元素是目前所有出现的颜色出现的次数。具体地，add 函数加入一个结点时，把它的颜色原来出现的次数从集合中删除，它的颜色出现的次数 +1，然后把现在的这个次数放入集合；del 函数也是类似的。一个区间解决完后，如果集合中最小的数与最大的数相等，就说明这个结点为根的子树是“颜色平衡树”，答案加上 1 即可。

需要注意的是，如果把块长设置为 \sqrt{n}，那么将会 TLE。考虑使用常数块长，经测试，块长为 6000 时可以通过本题。

放代码：

#pragma GCC target("avx,avx2") // 指令集优化
#include<bits/stdc++.h>
using namespace std;
typedef pair<int,int> pii;
int f[200001],c[200001],l[200001],t[200001],m[200001],o;
vector<int> g[200001];
vector<pii> q;
multiset<int> w; // 使用 STL 多重集合 multiset
inline int read(){
  int x=0; char c=getchar();
  while(!isdigit(c))c=getchar();
  while(isdigit(c))x=(x<<1)+(x<<3)+(c^48),c=getchar();
  return x;
} // 快速读入，用于卡常
void add(int x){
  if(t[c[x]])w.erase(w.find(t[c[x]]));
  w.emplace(++t[c[x]]);
}
void del(int x){
  w.erase(w.find(t[c[x]]));
  if(--t[c[x]])w.emplace(t[c[x]]);
}
void dfs(int u){
  m[l[u]=++o]=u;
  for(int i:g[u])dfs(i);
  q.emplace_back(l[u],o); // 该结点管辖的区间
}
int main(){
  int n=read(),s=0,L=1,R=0,sq=6000;
  for(int i=1;i<=n;i++)
    c[i]=read(),g[f[i]=read()].emplace_back(i);
  dfs(1); // 预处理
  sort(q.begin(),q.end(),[&](pii a,pii b){
    if(a.first/sq!=b.first/sq)return a.first<b.first;
    return a.first/sq&1?a.second>b.second:a.second<b.second;
    }); // 奇偶性排序
  for(auto [l,r]:q){
    while(L>l)add(m[--L]);
    while(R<r)add(m[++R]);
    while(L<l)del(m[L++]);
    while(R>r)del(m[R--]);
    s+=*w.begin()==*prev(w.end());
  } // 莫队算法
  printf("%d\n",s);
  return 0;
}