Codeforces 734E 树的直径

题意

传送门 Codeforces 734E Anton and Tree

题解

相同颜色的连通分量会被同时染色，故将其缩点，得到一个满足二分图性质的树。设树的直径为 $d$，在任一个节点染色，至多使树的直径减少 $2$，故染色数量的下界为 $\lceil d/2\rceil$。可以取到这个下界，取节点 $u$ 满足 $u$ 至树上任一节点距离不超过 $\lceil d/2\rceil$，那么进行 $\lceil d/2\rceil$ 次 $paint(u)$ 即可。

#include <bits/stdc++.h>
using namespace std;
#define rep(i, l, r) for (int i = l, _ = r; i < _; ++i)
#define pb push_back
typedef long long ll;
const int MAXN = 2E5 + 5;
int N, col[MAXN], cc[MAXN];
vector<int> G[MAXN], nG[MAXN];
int dp[MAXN], res;

void dfs(int u, int c, int k)
{
    cc[u] = k;
    for (auto &v : G[u])
        if (cc[v] == -1 && col[v] == c)
            dfs(v, c, k);
}

void get(int u, int p)
{
    for (auto &v : nG[u])
        if (v != p)
        {
            get(v, u);
            res = max(res, dp[u] + dp[v] + 1);
            dp[u] = max(dp[u], dp[v] + 1);
        }
}

int main()
{
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    cin >> N;
    rep(i, 0, N) cin >> col[i];
    rep(i, 1, N)
    {
        int u, v;
        cin >> u >> v;
        --u, --v;
        G[u].pb(v), G[v].pb(u);
    }
    memset(cc, -1, sizeof(cc));
    int k = 0;
    rep(i, 0, N)
    {
        if (cc[i] == -1)
            dfs(i, col[i], k++);
    }
    rep(u, 0, N) for (auto &v : G[u]) if (col[u] != col[v]) nG[cc[u]].pb(cc[v]);
    get(0, -1);
    res = (res >> 1) + (res & 1);
    cout << res << '\n';
    return 0;
}