P7215 [JOISC2020] 首都

本题大意就是合并最小的次数，使得整个这些城市的点相互联通。
考虑一个点分治，把分治中心作为首都城市，那么有一个贪心的思想，若是要合并的点不在当前分治树里就直接推出，因为若出现这种情况，那分治中心应该在上面。所以我们把分治中心的颜色的所有点拉出来，然后跑一个$ bfs $一样的东西。
就是把这些点的$ fa $的点也放进去，然后判断是否在子树里，颜色数$ +1 $最后如果做完了，就说明此时的分治中心是
某个最优中心，计算答案。

 

代码：

#include<bits/stdc++.h>
using namespace std;
int n,k;
int c[200005];
int siz[200005],vis[200005],fa[200005],root,ans=2147483647,num;
vector<int>p[200005],col[200005];
int zhan[200005],top=0,ok[200005],vist[200005];
queue<int>q;
void get_siz(int x,int ff){
  siz[x]=1;
  for(int j=0;j<p[x].size();j++){
    int to=p[x][j];
    if(to==ff||vis[to])continue;
    get_siz(to,x);
    siz[x]+=siz[to];
  }
}
void get_weight(int x,int ff,int gs){
  int mx=0;
  for(int j=0;j<p[x].size();j++){
    int to=p[x][j];
    if(to==ff||vis[to])continue;
    get_weight(to,x,gs);
    mx=max(mx,siz[to]);
  }
  mx=max(mx,gs-siz[x]);
  if(mx*2<=gs)root=x;
}
void dfs(int x,int ff){
  fa[x]=ff;
  zhan[++top]=x;
  ok[x]=1;
  for(int j=0;j<p[x].size();j++){
    int to=p[x][j];
    if(vis[to]||to==ff)continue;
    dfs(to,x);
  }
}
void solve(int x){
  num=0;
  while(!q.empty())q.pop();
  bool ye=0;
  vist[c[x]]=1;
  for(int j=0;j<col[c[x]].size();j++){
    int to=col[c[x]][j];
    if(ok[to]==0){
      ye=1;break;
    }
    q.push(to);
  }num++;
  if(ye)return;
  while(!q.empty()){
    int fi=q.front();q.pop();
    if(vist[c[fa[fi]]]==0){
      vist[c[fa[fi]]]=1;
      ye=0;
      for(int j=0;j<col[c[fa[fi]]].size();j++){
        int to=col[c[fa[fi]]][j];
        if(ok[to]==0){
          ye=1;break;
        }
        q.push(to);
      }
      num++;
      if(ye)return;
    }
  }
  ans=min(ans,num-1);
}
void dfz(int x){
  get_siz(x,x);
  root=0;
  get_weight(x,x,siz[x]);
  x=root;
  dfs(x,x);
  solve(x);
  while(top){
    ok[zhan[top]]=0;
    vist[c[zhan[top]]]=0;
    top--;
  }
  vis[x]=1;
  for(int j=0;j<p[x].size();j++){
    int to=p[x][j];
    if(vis[to])continue;
    dfz(to);
  }
}
int main(){
  // freopen("3.in","r",stdin);
  // freopen("3.out","w",stdout);
  scanf("%d%d",&n,&k);
  for(int i=1;i<n;i++){
    int n1,n2;
    scanf("%d%d",&n1,&n2);
    p[n1].push_back(n2);
    p[n2].push_back(n1);
  }
  for(int i=1;i<=n;i++){
    scanf("%d",&c[i]);
    col[c[i]].push_back(i);
  }
  dfz(1);
  printf("%d\n",ans);
  return 0;
}