YbtOJ 714「点分治」染色计划

#pragma GCC optimize("Ofast")
#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,avx2,fma")
#pragma GCC optimize("unroll-loops")
#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define W while
#define I inline
#define RI register int
#define LL long long
#define Cn const
#define CI Cn int&
#define gc getchar
#define D isdigit(c=gc())
#define pc(c) putchar((c))
#define min(x,y) ((x)<(y)?(x):(y))
#define max(x,y) ((x)>(y)?(x):(y))
using namespace std;
namespace Debug{
    Tp I void _debug(Cn char* f,Ty t){cerr<<f<<'='<<t<<endl;}
    Ts I void _debug(Cn char* f,Ty x,Ar... y){W(*f!=',') cerr<<*f++;cerr<<'='<<x<<",";_debug(f+1,y...);}
    Tp ostream& operator<<(ostream& os,Cn vector<Ty>& V){os<<"[";for(Cn auto& vv:V) os<<vv<<",";os<<"]";return os;}
    #define gdb(...) _debug(#__VA_ARGS__,__VA_ARGS__)
}using namespace Debug;
namespace FastIO{
    Tp I void read(Ty& x){char c;int f=1;x=0;W(!D) f=c^'-'?1:-1;W(x=(x<<3)+(x<<1)+(c&15),D);x*=f;}
    Ts I void read(Ty& x,Ar&... y){read(x),read(y...);}
    Tp I void write(Ty x){x<0&&(pc('-'),x=-x,0),x<10?(pc(x+'0'),0):(write(x/10),pc(x%10+'0'),0);}
    Tp I void writeln(Cn Ty& x){write(x),pc('\n');}
}using namespace FastIO;
Cn int N=2e5+10,inf=2e9;
int n,m,F[N],c[N],Ans=inf,used[N],Min,cnt,id,sz[N],vis[N];
#define PA pair<int,int>
#define MP make_pair
#define fi first
#define se second
vector<int> G[N],cg[N],vc[N],v;
#define pb push_back
I void GetG(CI x,CI fa){
    RI mx=0;sz[x]=1;for(auto i:G[x]) i^fa&&!vis[i]&&(GetG(i,x),sz[x]+=sz[i],mx=max(mx,sz[i]));
    mx=max(mx,cnt-sz[x]),mx<Min&&(Min=mx,id=x);
}
I void Get(CI x,CI fa){F[x]=fa,v.pb(x);for(auto i:G[x]) i^fa&&!vis[i]&&(Get(i,x),0);}
I void Cle(CI x,CI fa){F[x]=0;for(auto i:G[x]) i^fa&&!vis[i]&&(Cle(i,x),0);}
queue<int> q;
I void Dfs(CI x){
    RI i,j,X=0;vis[x]=1;v.clear(),Get(x,0);for(auto i:v) vc[c[i]].pb(i),used[c[i]]=0;
    W(!q.empty()) q.pop();used[c[x]]=1,q.push(c[x]);W(!q.empty()){
        RI u=q.front();q.pop(),X++;if(cg[u].size()!=vc[u].size()){X=inf;break ;}
        for(auto i:vc[u]) F[i]&&(!used[c[F[i]]]&&(q.push(c[F[i]]),used[c[F[i]]]=1));
    }Ans=min(Ans,X);for(auto i:v) vc[c[i]].clear();Cle(x,0);
    for(auto i:G[x]) !vis[i]&&(cnt=sz[i],Min=inf,GetG(i,0),Dfs(id),0);
}
int main(){
    freopen("color.in","r",stdin),freopen("color.out","w",stdout);
    RI i,x,y;for(read(n,m),i=1;i<n;i++) read(x,y),G[x].pb(y),G[y].pb(x);
    for(i=1;i<=n;i++) read(c[i]),cg[c[i]].pb(i);Min=inf,cnt=n,GetG(1,0),Dfs(id);
    return writeln(Ans-1),0;
}