Problem Description

You have an undirected tree with n nodes labeled from 0 to n - 1, connected by n - 1 edges. The tree structure is given as a 2D array edges, where each edges[i] = [ai, bi] represents an edge between nodes ai and bi.

Some nodes in the tree are marked as restricted, given in the array restricted. These restricted nodes cannot be visited.

Your task is to find the maximum number of nodes you can reach starting from node 0, without visiting any restricted nodes. Node 0 itself is guaranteed to never be restricted.

For example, if you have a tree where node 0 connects to nodes 1, 2, and 3, but node 2 is restricted, you can only reach nodes 0, 1, and 3 (assuming nodes 1 and 3 don't lead to other unrestricted nodes). The answer would be 3.

The solution uses depth-first search (DFS) to traverse the tree from node 0, counting all reachable nodes while avoiding the restricted ones. It builds an adjacency list to represent the tree structure and uses a visited set (initialized with restricted nodes) to track which nodes should not be explored.

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 nodes and edges. 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 exact definition of a tree structure.

DFS

• Conclusion: Since we're dealing with a tree structure, the flowchart leads us directly to DFS (Depth-First Search).

This makes perfect sense for our problem because:

1. We need to traverse from node 0 and explore all reachable nodes
2. We need to count nodes while avoiding restricted ones
3. DFS naturally explores all paths from a starting node in a tree
4. The recursive nature of DFS makes it easy to count nodes as we traverse

The flowchart confirms that DFS is the appropriate algorithm for this tree traversal problem where we need to explore all reachable nodes from a starting point while respecting certain constraints (restricted nodes).

Intuition

Since we've identified this as a tree traversal problem using DFS, let's think about what we're actually trying to do. We want to count all nodes we can reach from node 0, but we need to avoid the restricted nodes - they act like barriers that block our path.

Think of it like exploring a network of rooms where some doors are locked (restricted nodes). Starting from room 0, we want to count how many rooms we can actually visit. If we encounter a locked door, we simply don't go through it, which means we can't reach any rooms that are only accessible through that locked door.

The key insight is that restricted nodes effectively "cut off" portions of the tree. When we hit a restricted node during our traversal, it's like hitting a dead end - we can't go further in that direction. This naturally suggests a DFS approach where we:

1. Start from node 0
2. Explore each neighbor recursively
3. Skip any neighbor that is restricted
4. Count each node we successfully visit

To implement this efficiently, we can treat restricted nodes as if they're already visited before we even start our traversal. By adding them to our visited set initially, we ensure our DFS will never attempt to explore them. This is clever because it unifies our logic - we don't need separate checks for "is this restricted?" and "have I visited this?" - both cases are handled by the same visited set.

The recursive DFS naturally accumulates the count: each call returns 1 (for the current node) plus the sum of all nodes reachable from its unvisited neighbors. This builds up our total count as the recursion unwinds, giving us the final answer when we return from the initial call to dfs(0).

Learn more about Tree, Depth-First Search, Breadth-First Search, Union Find and Graph patterns.

Solution Approach

Let's implement the DFS solution step by step:

Step 1: Build the Graph Representation

First, we need to convert the edge list into an adjacency list for easier traversal. We use a defaultdict(list) to store the graph:

g = defaultdict(list)
for a, b in edges:
    g[a].append(b)
    g[b].append(a)

Since the tree is undirected, we add both directions for each edge. This allows us to traverse from any node to its neighbors.

Step 2: Initialize the Visited Set

Here's the clever part - we initialize our vis set with the restricted nodes:

vis = set(restricted)

By marking restricted nodes as "visited" from the start, we ensure our DFS will never explore them. This elegantly handles the restriction without needing additional conditional checks during traversal.

Step 3: Implement the DFS Function

The recursive DFS function counts reachable nodes:

def dfs(i: int) -> int:
    vis.add(i)
    return 1 + sum(j not in vis and dfs(j) for j in g[i])

Let's break down what happens in each call:

• vis.add(i): Mark the current node as visited to avoid cycles
• 1: Count the current node itself
• sum(j not in vis and dfs(j) for j in g[i]): For each neighbor j:
  • If j is not visited (and not restricted, since restricted nodes are already in vis)
  • Recursively call dfs(j) and add its count
  • The expression j not in vis and dfs(j) uses short-circuit evaluation - if j is already visited, it returns False (which equals 0 in the sum)

Step 4: Start the Traversal

Finally, we start the DFS from node 0:

return dfs(0)

The function will explore all reachable nodes from node 0, automatically avoiding restricted nodes, and return the total count.

Time Complexity: O(n) where n is the number of nodes, as we visit each reachable node exactly once.

Space Complexity: O(n) for the adjacency list, visited set, and recursion stack in the worst case.

Example Walkthrough

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

Example:

• Tree with 7 nodes (0-6)
• Edges: [[0,1], [1,2], [1,3], [3,4], [0,5], [5,6]]
• Restricted nodes: [2, 4]

The tree looks like this:

        0
       / \
      1   5
     / \   \
    2   3   6
        |
        4

Step 1: Build the adjacency list

g = {
    0: [1, 5],
    1: [0, 2, 3],
    2: [1],
    3: [1, 4],
    4: [3],
    5: [0, 6],
    6: [5]
}

Step 2: Initialize visited set with restricted nodes

vis = {2, 4}  # Nodes 2 and 4 are "pre-blocked"

Step 3: Execute DFS from node 0

Let's trace through the DFS execution:

1. dfs(0):

   • Add 0 to vis: vis = {2, 4, 0}
   • Check neighbors [1, 5]:
     • 1 not in vis → call dfs(1)
     • 5 not in vis → call dfs(5)
   • Returns: 1 + dfs(1) + dfs(5)

2. dfs(1):

   • Add 1 to vis: vis = {2, 4, 0, 1}
   • Check neighbors [0, 2, 3]:
     • 0 in vis → skip (returns 0)
     • 2 in vis (restricted) → skip (returns 0)
     • 3 not in vis → call dfs(3)
   • Returns: 1 + 0 + 0 + dfs(3) = 1 + dfs(3)

3. dfs(3):

   • Add 3 to vis: vis = {2, 4, 0, 1, 3}
   • Check neighbors [1, 4]:
     • 1 in vis → skip (returns 0)
     • 4 in vis (restricted) → skip (returns 0)
   • Returns: 1 + 0 + 0 = 1

4. Back to dfs(1): Returns 1 + 1 = 2

5. dfs(5):

   • Add 5 to vis: vis = {2, 4, 0, 1, 3, 5}
   • Check neighbors [0, 6]:
     • 0 in vis → skip (returns 0)
     • 6 not in vis → call dfs(6)
   • Returns: 1 + 0 + dfs(6)

6. dfs(6):

   • Add 6 to vis: vis = {2, 4, 0, 1, 3, 5, 6}
   • Check neighbors [5]:
     • 5 in vis → skip (returns 0)
   • Returns: 1 + 0 = 1

7. Back to dfs(5): Returns 1 + 1 = 2

8. Back to dfs(0): Returns 1 + 2 + 2 = 5

Final Answer: 5

We successfully reached nodes {0, 1, 3, 5, 6}, which is 5 nodes total. Nodes 2 and 4 were restricted and blocked our traversal, preventing us from exploring beyond them.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n)

The algorithm performs a depth-first search (DFS) starting from node 0. In the worst case, it visits all reachable nodes exactly once. Since each node is visited at most once (tracked by the vis set), and for each visited node we iterate through its adjacent edges, the total time complexity is O(V + E) where V is the number of vertices and E is the number of edges. Given that this is a tree structure (as implied by having n nodes and n-1 edges), we have E = n - 1. Therefore, the time complexity simplifies to O(n + (n-1)) = O(n).

Space Complexity: O(n)

The space complexity consists of several components:

• The adjacency list g stores all edges, which requires O(n) space since there are n-1 edges in a tree, and each edge is stored twice (bidirectional).
• The vis set can contain up to n nodes in the worst case, requiring O(n) space.
• The recursive DFS call stack can go up to depth n in the worst case (when the tree is a linear chain), requiring O(n) space.

Therefore, the overall space complexity is O(n).

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

Common Pitfalls

Pitfall 1: Forgetting to Handle Node 0 in Restricted Set

While the problem statement guarantees that node 0 is never restricted, some developers might accidentally add error-checking code that could break the solution:

Incorrect Approach:

def reachableNodes(self, n, edges, restricted):
    # Wrong: Trying to "handle" node 0 being restricted
    if 0 in restricted:
        return 0  # This will never happen per problem constraints
  
    # ... rest of the code

Why it's problematic: Adding unnecessary checks can introduce bugs and confusion. Trust the problem constraints.

Pitfall 2: Using Global Variables Incorrectly

When using the compact DFS with closure, forgetting that vis needs to be accessible can cause issues:

Incorrect Approach:

def reachableNodes(self, n, edges, restricted):
    g = defaultdict(list)
    for a, b in edges:
        g[a].append(b)
        g[b].append(a)
  
    def dfs(i):
        vis = set(restricted)  # Wrong: Creating new set in each call
        vis.add(i)
        return 1 + sum(j not in vis and dfs(j) for j in g[i])
  
    return dfs(0)

Correct Approach:

def reachableNodes(self, n, edges, restricted):
    g = defaultdict(list)
    for a, b in edges:
        g[a].append(b)
        g[b].append(a)
  
    vis = set(restricted)  # Correct: vis is in outer scope
  
    def dfs(i):
        vis.add(i)
        return 1 + sum(j not in vis and dfs(j) for j in g[i])
  
    return dfs(0)

Pitfall 3: Short-Circuit Evaluation Order

The compact expression j not in vis and dfs(j) relies on short-circuit evaluation. Reversing the order breaks the logic:

Incorrect Approach:

def dfs(i):
    vis.add(i)
    return 1 + sum(dfs(j) and j not in vis for j in g[i])  # Wrong order!

Why it fails: This calls dfs(j) first, which adds j to vis, making j not in vis always False. This leads to infinite recursion or incorrect results.

Pitfall 4: Modifying Restricted List

Attempting to modify the input restricted list directly instead of creating a new set:

Incorrect Approach:

def reachableNodes(self, n, edges, restricted):
    g = defaultdict(list)
    for a, b in edges:
        g[a].append(b)
        g[b].append(a)
  
    restricted.append(0)  # Wrong: Modifying input
    vis = restricted  # Wrong: vis is a list, not a set
  
    def dfs(i):
        if i not in vis:  # O(n) lookup in list!
            vis.append(i)
            # ... rest of logic

Problems:

1. Modifies the input array (side effects)
2. List lookups are O(n) instead of O(1)
3. Can cause duplicate entries in the list

Correct Approach: Always create a new set from the restricted list:

vis = set(restricted)  # Creates a new set, O(1) lookups

Pitfall 5: Handling Empty Graphs or Edge Cases

Not considering edge cases like when there are no edges or only restricted nodes:

Incomplete Approach:

def reachableNodes(self, n, edges, restricted):
    if not edges:  # What if n=1 and no edges?
        return 0  # Wrong: Should return 1 (node 0 itself)
  
    # ... rest of code

Correct Understanding:

• If n=1 and edges=[], the answer is 1 (just node 0)
• If all nodes except 0 are restricted, the answer is 1
• The base case is handled naturally by the DFS returning 1 for the starting node

The beauty of the provided solution is that it handles all these edge cases naturally without special conditionals.