LOJ#3280. 「JOISC 2020 Day4」首都城市 点分治+BFS

写代码时间：30-40min，调试时间：5min

我们发现，如果一个城市的一个点被选，则该城市其他点也都必须被选，可以考虑用点分治来解.

假设当前分治到的重心为 $x$，则只需考虑必经 $x$ 的连通块即可.

我们可以维护一个队列，开始的时候将重心的颜色放入，然后对每一种颜色的所有节点进行扩展：向上跳父亲，直到被访问过.

如果队列为空且不存在队列中一种颜色不在当前分治树块内，则可以更新一下答案.

总时间复杂度为 $O(n \log n)$ .

code: 

#include <bits/stdc++.h>  
#define N 200009  
#define ll long long 
#define setIO(s) freopen(s".in","r",stdin) 
using namespace std;     
vector<int>cn[N];  
queue<int>q;  
int n,K,root,sn,ans;  
int hd[N],to[N<<1],nex[N<<1],edges;   
int col[N],vis[N],size[N],mx[N],fa[N],cur[N],vised[N],solved[N];           
void add(int u,int v) 
{ 
    nex[++edges]=hd[u],hd[u]=edges,to[edges]=v;   
}
void getroot(int x,int ff) 
{
    size[x]=1,mx[x]=0;     
    for(int i=hd[x];i;i=nex[i]) 
    {
        int y=to[i]; 
        if(y==ff||vis[y]) continue;  
        getroot(y,x),size[x]+=size[y]; 
        mx[x]=max(mx[x],size[y]);     
    }   
    mx[x]=max(mx[x],sn-size[x]);  
    if(mx[x]<mx[root]) root=x; 
}         
void getsize(int x,int ff) 
{  
    size[x]=1;  
    for(int i=hd[x];i;i=nex[i]) 
        if(to[i]!=ff&&!vis[to[i]]) getsize(to[i],x),size[x]+=size[to[i]];  
}   
void getfa(int x,int ff) 
{
    fa[x]=ff,cur[col[x]]++;   
    for(int i=hd[x];i;i=nex[i]) if(to[i]!=ff&&!vis[to[i]]) getfa(to[i],x);   
}
// 当前重心 
void clr(int x,int ff) 
{  
    fa[x]=0,cur[col[x]]=0,solved[col[x]]=0,vised[x]=0;    
    for(int i=hd[x];i;i=nex[i]) if(to[i]!=ff&&!vis[to[i]]) clr(to[i],x);        
}
void check(int x) 
{      
    int flag=0,cnt=0; 
    getfa(x,0),q.push(col[x]);   
    while(!q.empty()) 
    {   
        int u=q.front(),p; q.pop();               
        if(solved[u]) 
            continue;  
        solved[u]=1,++cnt; 
        if(cur[u]!=cn[u].size()) {   flag=1; break;    }        
        for(int i=0;i<cn[u].size();++i) vised[cn[u][i]]=1;                 
        for(int i=0;i<cn[u].size();++i)     
            for(p=fa[cn[u][i]];p&&!vised[p];p=fa[p]) vised[p]=1,q.push(col[p]);     
    }     
    if(!flag) ans=min(ans,cnt);  
    while(!q.empty()) q.pop();   
    clr(x,0);   
}
void solve(int x) 
{      
    check(x);    
    vis[x]=1;   
    for(int i=hd[x];i;i=nex[i]) 
    {
        int y=to[i]; 
        if(vis[y]) continue;           
        getsize(y,x);  
        sn=size[y],root=0,getroot(y,x),solve(root);     
    }
}
int main() 
{ 
    // setIO("input");   
    int x,y,z; 
    scanf("%d%d",&n,&K),ans=n;    
    for(int i=1;i<n;++i) 
    { 
        scanf("%d%d",&x,&y);   
        add(x,y),add(y,x);  
    }   
    for(int i=1;i<=n;++i)    
        scanf("%d",&col[i]),cn[col[i]].push_back(i);         
    root=0,mx[0]=N,sn=n,getroot(1,0);      
    solve(root);          
    printf("%d\n",ans-1);   
    return 0; 
}