思路
把连通块缩成一个点,然后就可以直接转化成一个相邻点颜色不同的树。那么把这颗树变成同色,只需要(树的直径+1)/2就可以了。
很简单的思路,然年被我实现搞出了无数个问题。总之就是非常离谱
离谱的提交记录。
#include "bits/stdc++.h"
#define pii pair<int,int>
using namespace std;
const int N = 2e5 + 10;
int head[N << 1];
int head2[N << 1];
int cnt;
struct node {
int next, to;
} a[N << 1],b[N<<1];
int f[N];
int find(int x) {
return f[x] == x ? x : f[x] = find(f[x]);
}
void merge(int x, int y) {
x = find(x);
y = find(y);
if (x != y)
f[x] = y;
}
void add(int u, int v) {
a[++cnt].to = v;
a[cnt].next = head[u];
head[u] = cnt;
}
void add2(int u, int v) {
b[++cnt].to = v;
b[cnt].next = head2[u];
head2[u] = cnt;
}
bool col[N];
bool vis[N];
void dfs(int u, int fa) {
for (int i = head[u]; i; i = a[i].next) {
int v = a[i].to;
if (v != fa && col[v] == col[u] && !vis[v]) {
merge(v, u);
vis[v] = 1;
dfs(v, u);
}
}
}
int dp1[N];
int dp2[N];
int ans = 0;
vector<pii > st;
void dfs2(int u, int fa) {
for (int i = head2[u]; i; i = b[i].next) {
int v = b[i].to;
if (v == fa) continue;
dfs2(v, u);
if (dp1[u] < dp1[v] + 1) {
dp2[u] = dp1[u];
dp1[u] = dp1[v] + 1;
} else if (dp2[u] < dp1[v] + 1) {
dp2[u] = dp1[v] + 1;
}
}
ans = max(ans, dp1[u] + dp2[u]);
}
void solve() {
int n;
scanf("%d",&n);
for (int i = 1; i <= n; ++i) {
scanf("%d",&col[i]);
f[i] = i;
}
for (int i = 1; i <= n - 1; ++i) {
int u, v;
scanf("%d%d",&u,&v);
add(u, v);
add(v, u);
st.push_back({
u, v});
}
for (int i = 1; i <= n; ++i) {
if (!vis[i]) {
vis[i] = 1;
dfs(i, 0);
}
}
cnt = 0;
for (auto i : st) {
if (f[i.first] != f[i.second]) {
add2(f[i.first], f[i.second]);
add2(f[i.second], f[i.first]);
}
}
dfs2(f[1], 0);
printf("%d\n",(ans + 1) >> 1);
}
int main() {
// freopen("in.txt", "r", stdin);
solve();
}