codeforces 734E（DFS,树的直径（最长路））

 1 #include<bits/stdc++.h>
 2 using namespace std;
 3 typedef long long ll;
 4 const int N = 2e5 + 5;
 5 int col[N],refl[N],maxdis,maxpos;
 6 vector <int> G[N],E[N];
 7 bool vis[N];
 8 void dfs(int x,int cur)
 9 {
10     if(vis[x])
11         return;
12     vis[x] = 1;
13     if(col[x] != col[cur])
14     {
15         E[x].push_back(cur);
16         E[cur].push_back(x);
17         cur = x;
18     }
19     for(int i = 0; i < G[x].size(); i++)
20         dfs(G[x][i],cur);
21 }
22 void DFS(int x,int dis)
23 {
24     if(vis[x])
25         return;
26     vis[x] = 1;
27     if(dis > maxdis)
28     {
29         maxdis = dis;
30         maxpos = x;
31     }
32     for(int i = 0; i < E[x].size(); i++)
33         DFS(E[x][i],dis + 1);
34 }
35 int main()
36 {
37     int n;
38     scanf("%d",&n);
39     for(int i = 1; i <= n; i++)
40         scanf("%d",col + i);
41     for(int i = 1; i <= n - 1; i++)
42     {
43         int u,v;
44         scanf("%d %d",&u,&v);
45         G[u].push_back(v);
46         G[v].push_back(u);
47     }
48     dfs(1,1);
49     memset(vis,0,sizeof(vis));
50     DFS(1,0);
51     memset(vis,0,sizeof(vis));
52     DFS(maxpos,0);
53     printf("%d
",(maxdis + 1) >> 1);
54     return 0;
55 }