$dfs$缩点,树形$dp$。
首先将连通块缩点,缩点后形成一个黑白节点相间的树。接下来的任务就是寻找一个$root$,使这棵树以$root$为根,树的高度是最小的(也就是一层一层染色)。树形$dp$可以解决这个问题,第一次$dfs$处理子树,第二次$dfs$枚举$root$计算答案。
#pragma comment(linker, "/STACK:1024000000,1024000000") #include<cstdio> #include<cstring> #include<cmath> #include<algorithm> #include<vector> #include<map> #include<set> #include<queue> #include<stack> #include<ctime> #include<iostream> using namespace std; typedef long long LL; const double pi=acos(-1.0); void File() { freopen("D:\\in.txt","r",stdin); freopen("D:\\out.txt","w",stdout); } template <class T> inline void read(T &x) { char c = getchar(); x = 0; while(!isdigit(c)) c = getchar(); while(isdigit(c)) { x = x * 10 + c - '0'; c = getchar(); } } int n; int c[200010]; int nc[200010]; int u[200010]; int v[200010]; int belong[200010]; int block; vector<int>G[200010]; int flag[200010]; vector<int>L[200010],R[200010]; int ans; int a[200010]; void dfs(int x) { belong[x]=block; for(int i=0;i<G[x].size();i++) { int to=G[x][i]; if(c[to]!=c[x]) continue; if(belong[to]!=0) continue; dfs(to); } } void dp(int x) { flag[x]=1; int mx=0; for(int i=0;i<G[x].size();i++) { int to=G[x][i]; if(flag[to]) { L[x].push_back(mx); continue; } dp(to); mx=max(mx,a[to]); L[x].push_back(mx+1); } a[x]=mx+1; } void dp2(int x) { flag[x]=1; int mx=0; for(int i=G[x].size()-1;i>=0;i--) { int to=G[x][i]; if(flag[to]) { R[x].push_back(mx); continue; } dp2(to); mx=max(mx,a[to]); R[x].push_back(mx+1); } a[x]=mx+1; } void Find(int x,int h) { flag[x]=1; int nh; int sz=G[x].size(); for(int i=0;i<G[x].size();i++) { int root=G[x][i]; if(flag[root]) continue; nh=h+1; int k=0; if(i-1>=0) k=max(k,L[x][i-1]); if(sz-i-2>=0) k=max(k,R[x][sz-i-2]); if(k!=0) nh=max(nh,k+1); ans=min(ans,max(nh,a[root])); Find(root,nh); } } int main() { cin>>n; for(int i=1;i<=n;i++) cin>>c[i]; for(int i=1;i<=n-1;i++) { cin>>u[i]>>v[i]; G[u[i]].push_back(v[i]); G[v[i]].push_back(u[i]); } for(int i=1;i<=n;i++) { if(belong[i]!=0) continue; block++; nc[block]=c[i]; dfs(i); } for(int i=1;i<=n;i++) G[i].clear(); for(int i=1;i<=n-1;i++) { if(c[u[i]]==c[v[i]]) continue; G[belong[u[i]]].push_back(belong[v[i]]); G[belong[v[i]]].push_back(belong[u[i]]); } memset(flag,0,sizeof flag); dp(1); memset(flag,0,sizeof flag); dp2(1); ans=a[1]; memset(flag,0,sizeof flag); Find(1,1); printf("%d\n",ans-1); return 0; }