恐怕树上直径打错了(?)我家树上直径有点不同..

并茶集缩点+树上直径

#include <bits/stdc++.h>
#define maxn 200005
using namespace std;
int n, w[maxn], head[maxn], tot, ans1, ans2, dp[maxn], fa[maxn], fa2[maxn], faq[maxn], dep[maxn];
struct node {
	int v, pre;
}e[maxn << 1];
void add(int u, int v) {
	e[++tot].v = v;
	e[tot].pre = head[u];
	head[u] = tot;
}
int head2[maxn], tot2;
struct node1 {
	int v, pre;
}e2[maxn << 1];
void added(int u, int v) {
	e2[++tot2].v = v;
	e2[tot2].pre = head2[u];
	head2[u] = tot2;
}
void init() {
	for(int i = 1;i <= n;i++) faq[i] = i;
}
int get(int x) {
	if(x == faq[x]) return x;
	else return faq[x] = get(faq[x]);
}
void com(int u, int v) {
	int t1 = get(u), t2 = get(v);
	if(t1 != t2) faq[t1] = t2;
}
void dfs1(int now) {
	for(int i = head[now];i;i = e[i].pre) {
		int v = e[i].v;
		if(v != fa[now]) {
			fa[v] = now;
			if(w[v] == w[now]) com(v, now);
			dfs1(v);
			if(w[v] != w[now]) added(get(now), get(v)), added(get(v), get(now));
		}
	}
}
void dfs2(int now, int st) {
	bool flag = false;
	for(int i = head2[now];i;i = e2[i].pre) {
		int v = e2[i].v;
		if(v != fa2[now]) {
			flag = true;
			fa2[v] = now;
			dep[v] = dep[now] + 1;
			dfs2(v, st);
			dp[now] = max(dp[now], dp[v] + 1);
		}
	}
	if(!flag) dp[now] = 1;
	if(now == st) {
		for(int i = head2[now];i;i = e2[i].pre) {
			int v = e2[i].v;
			if(dp[v] > ans1) {
				ans2 = ans1;
				ans1 = dp[v];
			}
			else if(dp[v] > ans2 && dp[v] < ans1) {
				ans2 = dp[v];
			}
		}
	}
}
int main() {
	scanf("%d", &n);
	init();
	for(int i = 1;i <= n;i++) scanf("%d", &w[i]);
	for(int i = 1;i <= n - 1;i++) {
		int u, v;
		scanf("%d%d", &u, &v);
		add(u, v);
		add(v, u);
	}
	dfs1(1);
	for(int i = 1;i <= n;i++) {
		if(faq[i] == i) {
			dep[i] = 1;
			dfs2(i, i);
			break;
		}
	}
	cout << (ans1 + ans2 + 1) / 2 << endl;
}