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 = ;
    while(!isdigit(c)) c = getchar();
    while(isdigit(c))
    {
        x = x *  + c - '';
        c = getchar();
    }
}

int n;
int c[];
int nc[];
int u[];
int v[];
int belong[];
int block;
vector<int>G[];
int flag[];
vector<int>L[],R[];
int ans;

int a[];

void dfs(int x)
{
    belong[x]=block;
    for(int i=;i<G[x].size();i++)
    {
        int to=G[x][i];
        if(c[to]!=c[x]) continue;
        if(belong[to]!=) continue;
        dfs(to);
    }
}

void dp(int x)
{
    flag[x]=;

    int mx=;
    for(int i=;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+);
    }
    a[x]=mx+;
}

void dp2(int x)
{
    flag[x]=;

    int mx=;
    for(int i=G[x].size()-;i>=;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+);
    }
    a[x]=mx+;
}

void Find(int x,int h)
{
    flag[x]=;

    int nh;
    int sz=G[x].size();
    for(int i=;i<G[x].size();i++)
    {
        int root=G[x][i];
        if(flag[root]) continue;

        nh=h+;

        int k=;
        if(i->=) k=max(k,L[x][i-]);
        if(sz-i->=) k=max(k,R[x][sz-i-]);
        if(k!=) nh=max(nh,k+);

        ans=min(ans,max(nh,a[root]));
        Find(root,nh);
    }
}

int main()
{
    cin>>n;
    for(int i=;i<=n;i++) cin>>c[i];
    for(int i=;i<=n-;i++)
    {
        cin>>u[i]>>v[i];
        G[u[i]].push_back(v[i]);
        G[v[i]].push_back(u[i]);
    }

    for(int i=;i<=n;i++)
    {
        if(belong[i]!=) continue;
        block++; nc[block]=c[i]; dfs(i);
    }

    for(int i=;i<=n;i++) G[i].clear();

    for(int i=;i<=n-;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,,sizeof flag); dp();
    memset(flag,,sizeof flag); dp2();

    ans=a[];
    memset(flag,,sizeof flag); Find(,);

    printf("%d\n",ans-);

    return ;
}