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)}
bad = set()


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


def dfs(start, visited, want):
    visited.add(start)
    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:
        bad.add(i)

if not found:
    print("NO")
else:
    print("YES")
    print(ans)


I'd love some guidance on how I can make this faster -- thank you!
Some comments on the code:

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