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

Question

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 c_i.

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 c₁, c₂, ..., c_n (1 ≤ c_i ≤ 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!

Answer 1

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

Answer 2

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)