Problem Description

Alice has an undirected tree with n nodes labeled from 0 to n - 1. The tree is given as a 2D array edges where edges[i] = [ai, bi] indicates an edge between nodes ai and bi.

Alice wants Bob to find which node is the root of her tree. Bob makes several guesses about the parent-child relationships in the tree. Each guess consists of:

• Choosing two distinct integers u and v where an edge [u, v] exists in the tree
• Claiming that u is the parent of v (meaning u is closer to the root than v)

Bob's guesses are given as a 2D array guesses where guesses[j] = [uj, vj] means Bob guessed that uj is the parent of vj.

Alice tells Bob that at least k of his guesses are correct.

The task is to find how many different nodes could possibly be the root of Alice's tree, given that at least k of Bob's guesses must be true. Return 0 if no such tree exists.

For example, if Bob guesses that node 1 is the parent of node 2, this guess would be correct if:

• The root is node 1 (or any node where 1 is closer to the root than 2)
• The path from the root to node 2 passes through node 1

The key insight is that changing which node is the root changes which parent-child relationships are true in the tree. A guess [u, v] is correct for a particular root if, when traversing from that root, we reach u before v along the path to v.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: The problem explicitly involves an undirected tree with n nodes and edges connecting them. A tree is a special type of graph (connected and acyclic).

Is it a tree?

• Yes: The problem states we have an undirected tree with n nodes and n-1 edges, which is the definition of a tree structure.

DFS

• Yes: Since we're dealing with a tree structure, DFS (Depth-First Search) is the appropriate algorithm.

Conclusion: The flowchart leads us directly to using DFS for this tree problem.

Why DFS is the Right Choice

The problem requires us to:

1. Traverse the tree to count how many of Bob's guesses are correct when different nodes are chosen as the root
2. Efficiently re-root the tree and recalculate the count of correct guesses

DFS is ideal here because:

• We need to traverse the entire tree structure from different potential root nodes
• We need to maintain parent-child relationships as we traverse
• The re-rooting technique (also known as "rerooting DP" or "tree DP with root change") naturally uses DFS to:
  • First DFS: Calculate the initial count of correct guesses with node 0 as root
  • Second DFS: Dynamically adjust the count as we conceptually move the root to each node

The solution uses two DFS traversals:

• dfs1: Counts correct guesses when node 0 is the root
• dfs2: Reuses the count from the parent node and adjusts it based on edge direction changes when moving the root from parent to child

Intuition

The key insight is that when we change which node is the root of a tree, the parent-child relationships change throughout the tree. However, we don't need to rebuild the entire tree for each potential root - we can be clever about it.

Let's think about what happens when Bob makes a guess [u, v] (meaning u is the parent of v):

• This guess is correct only if, in the rooted tree, node u is actually on the path from the root to node v, and u is one step closer to the root than v.
• In other words, there must be a directed edge from u to v when we orient the tree edges away from the root.

Now, here's the clever part: instead of checking all n nodes as potential roots independently (which would be inefficient), we can use a technique called "rerooting" or "tree DP with root change."

Start by picking any node as the initial root (say node 0). Count how many of Bob's guesses are correct with this root. Let's call this count cnt.

When we move the root from a node i to its adjacent node j, only the edge between i and j changes direction:

• Before: i is the parent of j (edge goes from i to j)
• After: j is the parent of i (edge goes from j to i)

This means:

• If Bob guessed [i, j], this guess was correct before but incorrect after moving the root, so we subtract 1 from cnt
• If Bob guessed [j, i], this guess was incorrect before but correct after moving the root, so we add 1 to cnt

By doing a DFS traversal and updating cnt as we "move" the root through the tree, we can efficiently calculate the number of correct guesses for each possible root. We count how many nodes give us at least k correct guesses.

The formula for updating the count when moving root from i to j is: new_cnt = old_cnt - (guess[i,j] exists ? 1 : 0) + (guess[j,i] exists ? 1 : 0)

This approach runs in O(n) time since we visit each node once in each DFS, making it much more efficient than the naive O(n²) approach of rebuilding the tree for each potential root.

Learn more about Tree, Depth-First Search and Dynamic Programming patterns.

Solution Approach

The solution implements the tree DP with root change technique using two DFS traversals:

Step 1: Build the Graph and Process Guesses

• Convert the edge list to an adjacency list g where g[i] contains all nodes adjacent to node i
• Store all guesses in a hash map gs as (parent, child) pairs for O(1) lookup

g = defaultdict(list)
for a, b in edges:
    g[a].append(b)
    g[b].append(a)
gs = Counter((u, v) for u, v in guesses)

Step 2: First DFS - Count Correct Guesses with Node 0 as Root The dfs1 function traverses the tree with node 0 as the root and counts how many guesses are correct:

• For each edge (i, j) where i is the parent of j, check if this guess exists in gs
• If (i, j) is in gs, increment the counter cnt

def dfs1(i, fa):
    nonlocal cnt
    for j in g[i]:
        if j != fa:  # Avoid going back to parent
            cnt += gs[(i, j)]  # Check if this parent-child guess exists
            dfs1(j, i)  # Recursively process child with i as parent

Step 3: Second DFS - Reroot and Count Valid Roots The dfs2 function performs the rerooting technique:

• Start with the count from node 0 as root
• For each node, check if cnt >= k to determine if it's a valid root
• When moving the root from node i to its child j:
  • The edge (i, j) reverses to (j, i)
  • Subtract 1 if guess (i, j) exists (was correct, now incorrect)
  • Add 1 if guess (j, i) exists (was incorrect, now correct)

def dfs2(i, fa):
    nonlocal ans, cnt
    ans += cnt >= k  # Count this node if it has enough correct guesses
    for j in g[i]:
        if j != fa:
            # Adjust count for moving root from i to j
            cnt -= gs[(i, j)]  # Remove the i->j guess
            cnt += gs[(j, i)]  # Add the j->i guess
          
            dfs2(j, i)  # Process with j as potential root
          
            # Restore count when backtracking
            cnt -= gs[(j, i)]
            cnt += gs[(i, j)]

The restoration step is crucial - after exploring the subtree rooted at j, we need to restore cnt to its original value before processing the next child of i.

Time Complexity: O(n + m) where n is the number of nodes and m is the number of guesses Space Complexity: O(n + m) for storing the graph and guess hash map

The beauty of this approach is that we only need to track the changes in count as we conceptually move the root, rather than recalculating everything from scratch for each potential root.

Example Walkthrough

Let's walk through a small example to illustrate the solution approach.

Example Input:

• Tree edges: [[0,1], [1,2], [1,3]] (forms a tree with node 1 connected to 0, 2, and 3)
• Bob's guesses: [[1,2], [0,1], [3,1]] (Bob guesses 1→2, 0→1, and 3→1 are parent-child relationships)
• k = 2 (at least 2 guesses must be correct)

Tree Structure:

    0
    |
    1
   / \
  2   3

Step 1: Build Graph and Process Guesses

• Graph adjacency list: {0: [1], 1: [0,2,3], 2: [1], 3: [1]}
• Guess map: {(1,2): 1, (0,1): 1, (3,1): 1}

Step 2: First DFS with Root = 0 Starting from node 0, we traverse and count correct guesses:

• Visit 0 → 1: Check if (0,1) is a guess? Yes! cnt = 1
• From 1, visit 1 → 2: Check if (1,2) is a guess? Yes! cnt = 2
• From 1, visit 1 → 3: Check if (1,3) is a guess? No. cnt = 2

With node 0 as root, 2 guesses are correct: (0,1) and (1,2).

Step 3: Second DFS - Reroot and Check Each Node

Starting at node 0 with cnt = 2:

• Node 0: cnt = 2 >= 2 ✓ Valid root! ans = 1

Move root from 0 to 1:

• Edge (0,1) reverses to (1,0)
• Remove guess (0,1): cnt = 2 - 1 = 1
• Add guess (1,0): Not in guesses, so cnt = 1 + 0 = 1
• Node 1: cnt = 1 < 2 ✗ Not valid

Move root from 1 to 2:

• Edge (1,2) reverses to (2,1)
• Remove guess (1,2): cnt = 1 - 1 = 0
• Add guess (2,1): Not in guesses, so cnt = 0 + 0 = 0
• Node 2: cnt = 0 < 2 ✗ Not valid
• Backtrack: Restore cnt = 1

Move root from 1 to 3:

• Edge (1,3) reverses to (3,1)
• Remove guess (1,3): Not in guesses, so cnt = 1 - 0 = 1
• Add guess (3,1): Yes, it's a guess! cnt = 1 + 1 = 2
• Node 3: cnt = 2 >= 2 ✓ Valid root! ans = 2

Result: 2 nodes (0 and 3) can be valid roots.

Verification:

• With root = 0: Correct guesses are (0,1) and (1,2) = 2 correct ✓
• With root = 3: Correct guesses are (3,1) and (1,2) = 2 correct ✓
• With root = 1: Only (1,2) is correct = 1 correct ✗
• With root = 2: Only (0,1) is correct = 1 correct ✗

The rerooting technique efficiently found that nodes 0 and 3 satisfy the requirement of having at least k=2 correct guesses.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n + m)

The algorithm consists of building the graph and two DFS traversals:

• Building the adjacency list g from edges takes O(n) time where n is the number of edges
• Creating the Counter gs from guesses takes O(m) time where m is the number of guesses
• dfs1 traverses the entire tree once, visiting each node exactly once and checking each edge, taking O(n) time
• dfs2 also traverses the entire tree once, visiting each node exactly once. For each edge traversed, it performs constant time operations (updating cnt by checking presence in gs), taking O(n) time overall

Total time complexity: O(n) + O(m) + O(n) + O(n) = O(n + m)

Space Complexity: O(n + m)

The space usage includes:

• Adjacency list g stores all edges bidirectionally, using O(n) space
• Counter gs stores all guesses as key-value pairs, using O(m) space
• Recursion stack for DFS can go up to the height of the tree, which in worst case (skewed tree) is O(n)
• Other variables (cnt, ans) use O(1) space

Total space complexity: O(n) + O(m) + O(n) = O(n + m)

Learn more about how to find time and space complexity quickly.

Common Pitfalls

Pitfall 1: Forgetting to Restore State During Backtracking

The Problem: A critical mistake is forgetting to restore the correct_count after exploring each subtree in the rerooting phase. Without proper restoration, the count becomes corrupted and produces incorrect results for subsequent nodes.

Incorrect Code:

def reroot_and_count(node: int, parent: int) -> None:
    nonlocal valid_roots, correct_count
  
    if correct_count >= k:
        valid_roots += 1
  
    for neighbor in graph[node]:
        if neighbor != parent:
            # Update count before recursion
            correct_count -= guess_set.get((node, neighbor), 0)
            correct_count += guess_set.get((neighbor, node), 0)
          
            reroot_and_count(neighbor, node)
          
            # MISSING: Restoration of correct_count!

Why It Fails: Without restoration, the correct_count accumulates changes incorrectly. When processing the second child of a node, it uses the modified count from the first child rather than the original count for that node.

Solution: Always restore the state after recursion:

# Update count before recursion
correct_count -= guess_set.get((node, neighbor), 0)
correct_count += guess_set.get((neighbor, node), 0)

reroot_and_count(neighbor, node)

# Restore count after recursion
correct_count += guess_set.get((node, neighbor), 0)
correct_count -= guess_set.get((neighbor, node), 0)

Pitfall 2: Using Wrong Data Structure for Guess Lookup

The Problem: Using a list to check if a guess exists leads to O(m) lookup time for each edge, resulting in O(n*m) overall complexity.

Incorrect Code:

# Inefficient: Using list for guesses
guesses_list = guesses  # Direct use of input list

def count_correct_guesses(node: int, parent: int) -> None:
    nonlocal correct_count
    for neighbor in graph[node]:
        if neighbor != parent:
            # O(m) operation for each edge check!
            if [node, neighbor] in guesses_list:
                correct_count += 1

Solution: Use a set or Counter for O(1) lookup:

# Efficient: Convert to set/Counter
guess_set = Counter((parent, child) for parent, child in guesses)

# Now checking is O(1)
if (node, neighbor) in guess_set:
    correct_count += guess_set[(node, neighbor)]

Pitfall 3: Incorrect Edge Direction Handling

The Problem: Confusing which edge direction represents a correct guess when rerooting. Remember that a guess (u, v) means "u is the parent of v".

Incorrect Logic:

# Wrong: Adding when should subtract, or vice versa
correct_count += guess_set.get((node, neighbor), 0)  # Should be subtract!
correct_count -= guess_set.get((neighbor, node), 0)  # Should be add!

Why It's Wrong: When moving root from node to neighbor:

• The edge (node, neighbor) was a parent-child relationship (node was parent)
• After rerooting, it becomes child-parent (neighbor is now parent)
• If (node, neighbor) was a correct guess before, it's now incorrect
• If (neighbor, node) was an incorrect guess before, it's now correct

Solution: Always remember the transformation rule:

# When moving root from node to neighbor:
correct_count -= guess_set.get((node, neighbor), 0)  # Was correct, now wrong
correct_count += guess_set.get((neighbor, node), 0)  # Was wrong, now correct

Pitfall 4: Not Handling Duplicate Guesses

The Problem: If the same guess appears multiple times in the input, treating it as appearing only once will give wrong counts.

Incorrect Code:

# Using set instead of Counter loses duplicate information
guess_set = set((parent, child) for parent, child in guesses)

# This only adds 1 even if the guess appears multiple times
if (node, neighbor) in guess_set:
    correct_count += 1  # Should add the actual count!

Solution: Use Counter to preserve the count of each guess:

guess_set = Counter((parent, child) for parent, child in guesses)
correct_count += guess_set[(node, neighbor)]  # Adds the actual frequency