CF734E Anton and Tree

https://www.luogu.com.cn/problem/CF734E

哈啤题
首先缩点

考虑一条链，一定是选中点然后点$(len+1)/2$次

树就是直径长度

具体时限把边两边是不同颜色的边权设为1，然后找直径即可

code:

#include<bits/stdc++.h>
#define N 2000050
using namespace std;
struct edge {
    int v, nxt, c;
} e[N << 1];
int p[N], eid;
void init() {
    memset(p, -1, sizeof p);
    eid = 0;
}
void insert(int u, int v, int c) {
    e[eid].v = v;
    e[eid].c = c;
    e[eid].nxt = p[u];
    p[u] = eid ++;
}
int dis[N], vis[N], n, a[N];
queue<int> q;
int dfs(int S) {
    for(int i = 0; i <= n; i ++) dis[i] = 0, vis[i] = 0;
    while(q.size()) q.pop();
    dis[S] = 0; q.push(S);
    while(q.size()) {
        int u = q.front(); q.pop();
        if(vis[u]) continue;
        vis[u] = 1;
        for(int i = p[u]; i + 1; i = e[i].nxt) {
            int v = e[i].v;
            if(vis[v]) continue;
            dis[v] = dis[u] + e[i].c;
            q.push(v);
        }
    }
    int mx = 1;
    for(int i = 1; i <= n; i ++) if(dis[i] > dis[mx]) mx = i;
    return mx;
}
int main() {
    init();
    scanf("%d", &n);
    for(int i = 1; i <= n; i ++) scanf("%d", &a[i]);
    for(int i = 1; i < n; i ++) {
        int u, v;
        scanf("%d%d", &u, &v);
        insert(u, v, a[u] != a[v]), 
        insert(v, u, a[u] != a[v]);
    }
    int ans = dfs(dfs(1));
    printf("%d", (dis[ans] + 1) / 2);
    return 0;
}