CodeForces 734E Anton and Tree

$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;
}