习题：Anton and Tree（树的直径）

题目

传送门

思路

比较有意思的一种对树的直径的利用

我们考虑每一次染色的时候，最优方案的每一次染色一定会将一个颜色相同的联通块全部都翻转颜色，

接着我们考虑缩点，缩完点之后，连着的两个节点一定是颜色不一样的

我们考虑缩完点之后的数的直径，我们将直径的中心节点提起来当根，

这颗树深度一样的点的颜色一定是一样的

此时的最优策略一定有一种是一层一层地染

最后一次染的点一定是数的直径上的一个点

我们设直径为\(d\)

最后的答案即为\(\frac{d+1}{2}\)

代码

#include<iostream>
#include<cstring>
#include<vector>
#include<queue>
using namespace std;
namespace ufs
{
    int fa[200005];
    void makeset(int n)
    {
        for(int i=1;i<=n;i++)
            fa[i]=i;
    }
    int findset(int u)
    {
        if(fa[u]==u)
            return u;
        return fa[u]=findset(fa[u]);
    }
    void unionset(int a,int b)
    {
        int u=findset(a);
        int v=findset(b);
        fa[u]=v;
    }
}
using namespace ufs;
struct node
{
    int u;
    int step;
};
int n;
int col[200005];
vector<int> g[200005],tre[200005];
bool used[200005];
void dfs(int u,int fa)
{
    for(int i=0;i<g[u].size();i++)
    {
        int v=g[u][i];
        if(v!=fa)
        {
            if(col[v]==col[u])
                unionset(u,v);
            dfs(v,u);
        }
    }
}
node bfs(int st)
{
    used[st]=1;
    node ret;
    queue<node> q;
    q.push((node){st,1});
    while(!q.empty())
    {
        node t=q.front();
        ret=t;
        q.pop();
        for(int i=0;i<tre[t.u].size();i++)
        {
            int v=tre[t.u][i];
            if(used[v]==0)
            {
                q.push((node){v,t.step+1});
                used[v]=1;
            }
        }
    }
    return ret;
}
int main()
{
    ios::sync_with_stdio(false);
    cin>>n;
    makeset(n);
    for(int i=1;i<=n;i++)
        cin>>col[i];
    for(int i=1,u,v;i<n;i++)
    {
        cin>>u>>v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1,0);
    for(int i=1;i<=n;i++)
        for(int j=0;j<g[i].size();j++)
            tre[findset(i)].push_back(findset(g[i][j]));
    memset(used,0,sizeof(used));
    node s1=bfs(findset(1));
    memset(used,0,sizeof(used));
    node s2=bfs(findset(s1.u));
    cout<<s2.step/2;
    return 0;
}