Anton and Tree

Anton and Tree

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

DFS/BFS

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

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

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

def getAns(v):
    vis[v] = 1
    global ans
    m1, m2 = 0, 0
    for x in edge[v]:
        if vis[x] == 0:
            t = getAns(x)
            if t > m1:
                m1, m2 = t, m1
            elif t > m2:
                m2 = t
    ans = max(ans, m1 + m2 + 1)
    return m1 + 1

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

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

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

代码如下：

# import sys
# sys.setrecursionlimit(1000000000)
from collections import defaultdict, deque


def buildTree(root):
    q = deque()
    q.append((root, root))
    vis = [0 for _ in range(n + 1)]
    vis[root] = 1
    while q:
        v, fa = q.pop()
        for x in preedge[v]:
            if vis[x] != 1:
                vis[x] = 1
                if col[x] == col[fa]:
                    q.append((x, fa))
                else:
                    edge[fa].append(x)
                    edge[x].append(fa)
                    q.append((x, x))

def getDeep(root):
    vis = [0 for _ in range(n + 1)]
    d = {}
    q = deque()
    q.append(root)
    vis[root] = 1
    d[root] = 0
    while q:
        u = q.pop()
        for x in edge[u]:
            if vis[x] == 0:
                vis[x] = 1
                d[x] = d[u] + 1
                q.append(x)
    return d

n = int(input())
precol = [int(_) for _ in input().split()]
preedge = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    x, y = [int(_) for _ in input().split()]
    preedge[x].append(y)
    preedge[y].append(x)

root = (n + 1) // 2
edge = [[] for _ in range(n + 1)]
col = [0 for _ in range(n + 1)]
for x in range(1, n + 1):
    col[x] = precol[x - 1]
buildTree(root)

d = getDeep(root)
u = max(d.keys(),key=(lambda k:d[k]))
print((max(getDeep(u).values()) + 1 ) // 2)

由于Python栈空间过小的原因，一直在test15 RE，改了一天，感觉还是C/C++用起来舒心=。=