Anton and Tree-CSDN博客

Anton and Tree

题目链接：http://codeforces.com/contest/734/problem/E

DFS/BFS

每一次操作都可以使连通的结点变色，所以可以将连通的点缩成一个点（由于Python的栈空间过小，这个操作只能用bfs；其他语言如c++可以直接dfs）。

由于缩点后的树的结点是黑白相间的，将整棵树变成相同颜色的最小操作数为：在树的中点不断操作将整棵树变成相同颜色所需要的操作数，也就是整棵树的最大半径。

求树的最大半径可以用一次dfs求得：

 1 def getAns(v):
 2     vis[v] = 1
 3     global ans
 4     m1, m2 = 0, 0
 5     for x in edge[v]:
 6         if vis[x] == 0:
 7             t = getAns(x)
 8             if t > m1:
 9                 m1, m2 = t, m1
10             elif t > m2:
11                 m2 = t
12     ans = max(ans, m1 + m2 + 1)
13     return m1 + 1

View Code

然而Python栈空间过小，只能考虑其他做法，这里使用的是两个bfs：

第一个bfs找出以root为根的深度最大的结点u（树的直径两端的结点必有一个为此结点）

第二次bfs找出以结点u为根的深度最大的结点v，此时v的深度则为树的直径

代码如下：

 1 # import sys
 2 # sys.setrecursionlimit(1000000000)
 3 from collections import defaultdict, deque
 4 
 5 
 6 def buildTree(root):
 7     q = deque()
 8     q.append((root, root))
 9     vis = [0 for _ in range(n + 1)]
10     vis[root] = 1
11     while q:
12         v, fa = q.pop()
13         for x in preedge[v]:
14             if vis[x] != 1:
15                 vis[x] = 1
16                 if col[x] == col[fa]:
17                     q.append((x, fa))
18                 else:
19                     edge[fa].append(x)
20                     edge[x].append(fa)
21                     q.append((x, x))
22 
23 def getDeep(root):
24     vis = [0 for _ in range(n + 1)]
25     d = {}
26     q = deque()
27     q.append(root)
28     vis[root] = 1
29     d[root] = 0
30     while q:
31         u = q.pop()
32         for x in edge[u]:
33             if vis[x] == 0:
34                 vis[x] = 1
35                 d[x] = d[u] + 1
36                 q.append(x)
37     return d
38 
39 n = int(input())
40 precol = [int(_) for _ in input().split()]
41 preedge = [[] for _ in range(n + 1)]
42 for _ in range(n - 1):
43     x, y = [int(_) for _ in input().split()]
44     preedge[x].append(y)
45     preedge[y].append(x)
46 
47 root = (n + 1) // 2
48 edge = [[] for _ in range(n + 1)]
49 col = [0 for _ in range(n + 1)]
50 for x in range(1, n + 1):
51     col[x] = precol[x - 1]
52 buildTree(root)
53 
54 d = getDeep(root)
55 u = max(d.keys(),key=(lambda k:d[k]))
56 print((max(getDeep(u).values()) + 1 ) // 2)