We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsCompaniesDisclaimersPrivacyTermsContactAbout Us

Count Number of Possible Root Nodes

Number: 2652

Difficulty: Hard

Paid? No

Companies: N/A

Problem Description

Given an undirected tree with n nodes and a list of guesses (each guess suggests a parent-child relation along an edge), determine how many nodes can be chosen as the root such that at least k of Bob's guesses turn out to be correct.

Key Insights

  • The tree is an undirected structure but every choice of the root produces a directed (parent-child) tree.
  • Start by rooting the tree arbitrarily (e.g., at node 0) and compute the number of correct guesses.
  • Use a re-rooting technique while performing DFS to efficiently compute the correct guess count for every possible root by updating the count when moving the root from a parent to one of its children.
  • When re-rooting along an edge (u, v), adjust the correct guess count:
    • Subtract 1 if the guess (u, v) was counted (because u was parent of v in the original rooting).
    • Add 1 if the guess (v, u) becomes valid in the new rooting.
  • Use a set or hash for quick lookup of guess edges.

Space and Time Complexity

Time Complexity: O(n + g) where n is the number of nodes and g is the number of guesses.
Space Complexity: O(n + g) due to the adjacency list and the guess set.

Solution

The solution involves two main DFS traversals:

  1. The initial DFS from an arbitrary root (node 0) calculates the number of correct guesses based on the initial parent-child relationships.
  2. A re-root DFS traverses from the initial root to all other nodes. When moving from node u to its child v, the count of correct guesses is updated based on:
    • Removing the correct count if the guess (u, v) held in the old rooting.
    • Adding the correct count if the guess (v, u) holds in the new rooting. By counting how many nodes have a correct guess count greater than or equal to k, we answer the problem.

Data structures used:

  • An adjacency list for representing the undirected tree.
  • A set (or equivalent) for constant time lookups of guesses.
  • Recursion (DFS) for tree traversal.

This approach leverages the re-rooting trick to efficiently update the result for all possible roots in O(n) time.

Code Solutions

# Python solution
from collections import defaultdict
import sys
sys.setrecursionlimit(10**6)

def countPossibleRoots(edges, guesses, k):
    # Build tree as an adjacency list
    tree = defaultdict(list)
    for u, v in edges:
        tree[u].append(v)
        tree[v].append(u)
    
    # Build guess set for O(1) lookup
    guess_set = set((u, v) for u, v in guesses)
    
    n = len(tree)
    # Global result count
    result = 0
    
    # First DFS: count correct guesses when 0 is the root.
    def dfs_count(node, parent):
        count = 0
        for neighbor in tree[node]:
            if neighbor == parent:
                continue
            # If current edge (node -> neighbor) is guessed correctly.
            if (node, neighbor) in guess_set:
                count += 1
            # Add correct guesses from subtree.
            count += dfs_count(neighbor, node)
        return count

    initial_count = dfs_count(0, -1)
    
    # Second DFS: re-rooting DFS to count valid roots.
    def dfs_reroot(node, parent, current_correct):
        nonlocal result
        # Add to result if current correct count meets requirement.
        if current_correct >= k:
            result += 1
        for neighbor in tree[node]:
            if neighbor == parent:
                continue
            # Save the current count
            temp = current_correct
            # When re-rooting from node to neighbor:
            # if guess (node, neighbor) was counted, it is lost.
            if (node, neighbor) in guess_set:
                current_correct -= 1
            # if guess (neighbor, node) becomes valid, add it.
            if (neighbor, node) in guess_set:
                current_correct += 1
            dfs_reroot(neighbor, node, current_correct)
            # Backtrack to restore current_correct for other branches.
            current_correct = temp

    dfs_reroot(0, -1, initial_count)
    return result

# Example usage:
edges = [[0,1],[1,2],[1,3],[4,2]]
guesses = [[1,3],[0,1],[1,0],[2,4]]
k = 3
print(countPossibleRoots(edges, guesses, k))  # Expected output: 3
← Back to All Questions