# Finding a vertex of a tree such that every subtree rooted at that vertex is of a uniform color

I'm trying to optimally solve codeforces problem 763A:

Each New Year Timofey and his friends cut down a tree of n vertices and bring it home. After that they paint all the n its vertices, so that the i-th vertex gets color ci.

Now it's time for Timofey birthday, and his mother asked him to remove the tree. Timofey removes the tree in the following way: he takes some vertex in hands, while all the other vertices move down so that the tree becomes rooted at the chosen vertex. After that Timofey brings the tree to a trash can.

Timofey doesn't like it when many colors are mixing together. A subtree annoys him if there are vertices of different color in it. Timofey wants to find a vertex which he should take in hands so that there are no subtrees that annoy him. He doesn't consider the whole tree as a subtree since he can't see the color of the root vertex.

A subtree of some vertex is a subgraph containing that vertex and all its descendants.

Your task is to determine if there is a vertex, taking which in hands Timofey wouldn't be annoyed.

Input
The first line contains single integer n (2 ≤ n ≤ 10⁵) — the number of vertices in the tree.

Each of the next n - 1 lines contains two integers u and v (1 ≤ u, v ≤ n, u ≠ v), denoting there is an edge between vertices u and v. It is guaranteed that the given graph is a tree.

The next line contains n integers c1, c2, ..., cn (1 ≤ ci ≤ 10⁵), denoting the colors of the vertices.

Output
Print "NO" in a single line, if Timofey can't take the tree in such a way that it doesn't annoy him.

Otherwise print "YES" in the first line. In the second line print the index of the vertex which Timofey should take in hands. If there are multiple answers, print any of them.

I have come up with a solution that works but times out on large inputs.

from collections import defaultdict

n = int(input())
graph = defaultdict(list)
colors = defaultdict(lambda: -1)
for i in range(n-1):
(u, v) = [int(i) for i in input().split()]
graph[u].append(v)
graph[v].append(u)

c = [int(i) for i in input().strip().split()]
colors = {i+1: v for (i, v) in enumerate(c)}

def start(start):
nebs = graph[start]
return all(dfs(n, set([start]), colors[n]) for n in nebs)

def dfs(start, visited, want):
if colors[start] != want:
return False
for neb in graph[start]:
if neb not in visited:
if not dfs(neb, visited, want):
return False
return True

found = False
ans = -1

for i in range(1, n+1):
if start(i):
found = True
ans = i
break
else:

print("NO")
else:
print("YES")
print(ans)


I'd love some guidance on how I can make this faster -- thank you!

• @dfhwzw edited to reflect I'm trying to get an optimal solution – nz_21 Jul 2 at 21:00
• Any inclination to interpret while all the other vertices move down? The original root becomes one of the children of the new root, the edge to the new root removed? All edges from the original root to the new root get reversed? – greybeard Jul 4 at 8:15

for i in range(n-1):
(u, v) = [int(i) for i in input().split()]


You are using i as the outer loop index, and i as the inner (list comprehension) loop index. This is confusing, and should be discouraged.

The second line may be written more simply as:

    u, v = map(int, input().split())


The assignment:

colors = defaultdict(lambda: -1)


is overwritten by the code that creates the colors dictionary (below), and so can be removed.

You are changing the variable meaning in the following lines:

c = [int(i) for i in input().strip().split()]
colors = {i+1: v for (i, v) in enumerate(c)}


In the first line, i is a colour value; in the second line, v is the colour value and i is the index of the colour value. For consistency, I'd used v instead of i in the first line. Or get rid of the list comprehension index altogether:

c = map(int, input().strip().split())
colors = { i: v for i, v in enumerate(c, 1) }


Using enumerate(c, 1) eliminates the need to use i+1 as the key expression.

Is the .strip() necessary?

I'm not certain what the point of maintaining your visited set is. Each time you start() a search from a new point, you are resetting the the visited set, and during any dfs(), you won't encounter a node you've already visited because the graph is guaranteed to be be a tree. You just need to pass the start node for each dfs() step as a previous node (instead of a visited set) to prevent back-tracking.

Use better variable names. I just figured out the nebs is an abbreviation for neighbours.

Faster solution hint:

I don’t believe a DFS is necessary.

• Ignore all edges between vertices of the same colour
• All other edges must share a common vertex (or Timofey will be annoyed)
• Examine first two edges between different coloured vertices. Continue with a linear search of remaining edges
• you're absolutely correct on all counts! Thanks -- I'll make the edits – nz_21 Jul 2 at 21:36
• @nz_21: Please be aware that you are not allowed to change the original code in the question after receiving an answer. – AlexV Jul 2 at 21:40
• Ah, I see - understood! – nz_21 Jul 2 at 21:43

The approach outlined in the question is brute-force: for each vertex, we do a DFS on all its subtrees, checking if they're all monochromatic. O(N^2)

Instead of running dfs on every vertex, we can selectively choose vertices that would conclusively tell us if it was possible to split up the graph as desired.

Consider a path: 111112111111. It doesn't feel right to run dfs on each of the vertices. It makes more sense to run it on JUST 2, as it it colored differently than than the vertex it is connected to.

Concretely: If we see a an edge that connects two vertices of different color, run a dfs on both of the them.

• If either of them can split up the graph the way we want, all good - just return any of them.

• If none of them can, we're screwed: if we try to root the graph on any other vertex, one of its subtrees WILL have this pair of different colored vertices. Why is this guaranteed? Because it's a tree.

Thus, at most, you'll have to make 2 dfs() calls, while scanning through all edges.

O(N)