Problem Description

You are given an undirected tree with n nodes labeled from 0 to n - 1, rooted at node 0. The tree structure is defined by a 2D integer array edges of length n - 1, where each edges[i] = [aᵢ, bᵢ] represents an edge between nodes aᵢ and bᵢ.

A node is considered good if all the subtrees rooted at its children have the same size (same number of nodes).

Your task is to count and return the total number of good nodes in the tree.

Key Definitions:

• A subtree of a node consists of that node and all of its descendants
• A node with no children (leaf node) is considered good by default since it has no subtrees to compare
• The root node checks if all its children's subtrees have equal sizes

Example Understanding:

If a node has 3 children, and their subtrees contain 5, 5, and 5 nodes respectively, then this node is good. But if the subtrees contain 5, 3, and 5 nodes, then this node is not good because not all subtree sizes are equal.

The problem asks you to traverse the entire tree and count how many nodes satisfy this "good" condition.

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 mentions 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 clearly states we have an undirected tree with n nodes labeled from 0 to n - 1, rooted at node 0. The tree structure is defined by edges connecting nodes.

DFS

• Yes: Since we're dealing with a tree structure and need to:

  1. Calculate subtree sizes for each node's children
  2. Compare these subtree sizes to determine if a node is "good"
  3. Process each node and its descendants recursively

  DFS is the natural choice here because:

  • We need to explore each subtree completely to calculate its size
  • We need to process children before determining if a parent is "good"
  • The recursive nature of DFS perfectly matches the recursive structure of trees

Conclusion: The flowchart suggests using DFS (Depth-First Search) for this problem. This aligns perfectly with the solution approach, where we recursively traverse the tree, calculate subtree sizes bottom-up, and check at each node whether all its children's subtrees have equal sizes.

Intuition

To determine if a node is "good", we need to know the sizes of all subtrees rooted at its children. This immediately suggests a bottom-up approach where we first calculate information about children before processing the parent.

Think about it this way: if we're standing at a node and want to check if it's good, we need to look down at each child and ask "how many nodes are in your subtree?" If all children answer with the same number, our node is good.

This naturally leads to a recursive DFS approach because:

1. We need complete subtree information: To know if a node is good, we must fully explore each of its children's subtrees to count their nodes. We can't make a decision about a node until we've completely processed all its descendants.

2. Information flows bottom-up: Leaf nodes are trivially good (they have no children to compare). As we move up the tree, each node needs the subtree sizes from its children, which those children calculated when they examined their own children.

3. Each node does two jobs:

   • Count the total nodes in its subtree (including itself)
   • Check if all its children's subtrees are equal in size

The key insight is that while traversing, we can kill two birds with one stone: as we recursively calculate subtree sizes for the parent's use, we can also check at each node whether it qualifies as "good" by comparing the sizes we get from different children.

For example, if a node has three children returning subtree sizes of [5, 5, 5], we know it's good. But if they return [5, 3, 5], it's not good. We accumulate the count of good nodes as we traverse, and also return the total subtree size (1 + 5 + 5 + 5 = 16) for our parent to use.

Learn more about Tree and Depth-First Search patterns.

Solution Approach

The implementation uses DFS with an adjacency list representation of the tree. Let's walk through the key components:

1. Building the Graph Structure

First, we construct an adjacency list g from the edges array. Since the tree is undirected, we add both directions for each edge:

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

2. The DFS Function

The core logic is in dfs(a, fa) where:

• a is the current node being processed
• fa is the parent node (to avoid revisiting in an undirected tree)

The function returns the total number of nodes in the subtree rooted at a.

3. Key Variables in DFS

• pre = -1: Stores the size of the first child's subtree (initialized to -1 to detect the first child)
• cnt = 1: Counts total nodes in current subtree (starts at 1 for the current node)
• ok = 1: Flag indicating if current node is good (assumes good initially)

4. Processing Children

For each neighbor b of node a:

for b in g[a]:
    if b != fa:  # Skip the parent to avoid cycles
        cur = dfs(b, a)  # Get subtree size of child b
        cnt += cur       # Add to total count

5. Checking if Node is Good

The clever part is comparing subtree sizes:

if pre < 0:
    pre = cur  # First child: store its subtree size
elif pre != cur:
    ok = 0     # Different size than first child: not good

• For the first child, we store its subtree size in pre
• For subsequent children, we compare their subtree sizes with pre
• If any child has a different subtree size, we mark ok = 0

6. Accumulating the Answer

After processing all children:

nonlocal ans
ans += ok  # Add 1 if good, 0 if not good
return cnt # Return total subtree size for parent's use

7. Starting the Traversal

We start DFS from the root (node 0) with parent -1:

ans = 0
dfs(0, -1)
return ans

The algorithm efficiently combines two tasks in a single traversal: counting subtree sizes (needed for parent nodes) and determining good nodes (our actual goal). The time complexity is O(n) since we visit each node once, and space complexity is O(n) for the recursion stack and adjacency list.

Example Walkthrough

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

Consider a tree with 7 nodes and the following edges:

edges = [[0,1], [0,2], [1,3], [1,4], [2,5], [2,6]]

This creates the following tree structure:

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

Step-by-step DFS traversal:

1. Start at node 0 (root, parent = -1)

   • Visit child 1

2. At node 1 (parent = 0)

   • Visit child 3 (leaf): returns subtree size = 1
   • Store pre = 1 (first child's subtree size)
   • Visit child 4 (leaf): returns subtree size = 1
   • Compare: pre (1) == cur (1) ✓ Same size!
   • Node 1 is good (ok = 1), increment ans = 1
   • Return subtree size = 1 + 1 + 1 = 3 to parent

3. Back at node 0

   • First child (node 1) returned size = 3
   • Store pre = 3
   • Visit child 2

4. At node 2 (parent = 0)

   • Visit child 5 (leaf): returns subtree size = 1
   • Store pre = 1 (first child's subtree size)
   • Visit child 6 (leaf): returns subtree size = 1
   • Compare: pre (1) == cur (1) ✓ Same size!
   • Node 2 is good (ok = 1), increment ans = 2
   • Return subtree size = 1 + 1 + 1 = 3 to parent

5. Back at node 0 (continuing)

   • Second child (node 2) returned size = 3
   • Compare: pre (3) == cur (3) ✓ Same size!
   • Node 0 is good (ok = 1), increment ans = 3
   • Return subtree size = 1 + 3 + 3 = 7

6. Process leaf nodes (3, 4, 5, 6)

   • Each leaf has no children, so pre remains -1
   • All leaves are good by default
   • Final ans = 3 + 4 = 7

Result: All 7 nodes are good!

• Node 0 is good: both children have subtrees of size 3
• Nodes 1 and 2 are good: both have children with subtrees of size 1
• Nodes 3, 4, 5, 6 are good: they're leaves with no children to compare

Solution Implementation

Time and Space Complexity

Time Complexity: O(n)

The algorithm performs a depth-first search (DFS) traversal of the tree. Each node is visited exactly once during the traversal. For each node, we iterate through its adjacent nodes (edges). Since the input represents a tree with n nodes, there are exactly n-1 edges, and each edge is traversed twice (once from each direction). Therefore, the total number of operations is proportional to the number of nodes plus the number of edges, which gives us O(n + (n-1)) = O(n).

Space Complexity: O(n)

The space complexity consists of:

1. The adjacency list g which stores all edges. For a tree with n nodes and n-1 edges, each edge appears twice in the adjacency list (once for each direction), requiring O(2(n-1)) = O(n) space.
2. The recursion stack for DFS, which in the worst case (a skewed tree) can go up to depth n, requiring O(n) space.
3. Additional variables like pre, cnt, ok, and ans use constant space O(1).

The overall space complexity is therefore O(n).

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

Common Pitfalls

1. Forgetting to Handle Leaf Nodes Correctly

A common mistake is not recognizing that leaf nodes (nodes with no children) should automatically be considered "good" nodes. The code might incorrectly skip counting them if the logic isn't properly structured.

Pitfall Example:

def dfs(current_node, parent_node):
    children = []
    for child in adjacency_list[current_node]:
        if child != parent_node:
            children.append(child)
  
    # Wrong: Leaf nodes won't be counted as good
    if len(children) == 0:
        return 1  # Only returns size, doesn't count as good node

Solution: Initialize is_good_node = True at the start, which naturally handles leaf nodes since they have no children to compare, maintaining their "good" status.

2. Using Global Variables Without Proper Scoping

When using a counter variable across recursive calls, forgetting to use nonlocal or class attributes can lead to incorrect results.

Pitfall Example:

def countGoodNodes(self, edges):
    # ...
    ans = 0  # Local variable
  
    def dfs(current_node, parent_node):
        # ...
        ans += is_good_node  # Error: UnboundLocalError
        return subtree_size

Solution: Either use nonlocal ans inside the DFS function or use a class attribute like self.good_nodes_count to maintain the counter across recursive calls.

3. Incorrect Parent Tracking in Undirected Tree

Since the tree is undirected, failing to track the parent node properly can cause infinite recursion or incorrect traversal.

Pitfall Example:

def dfs(current_node):  # Missing parent parameter
    for child in adjacency_list[current_node]:
        # Will revisit parent, causing infinite recursion
        child_size = dfs(child)

Solution: Always pass and check the parent node to avoid traversing back up the tree:

def dfs(current_node, parent_node):
    for child in adjacency_list[current_node]:
        if child != parent_node:  # Skip parent
            child_size = dfs(child, current_node)

4. Edge Case: Single Node Tree

When the tree has only one node (n=1), the edges array is empty. Not handling this case can cause errors when building the adjacency list or starting DFS.

Pitfall Example:

def countGoodNodes(self, edges):
    if not edges:  # Might return wrong answer
        return 0   # Wrong: single node is good

Solution: The current implementation handles this correctly because:

• An empty edges array creates an empty adjacency list
• DFS on node 0 with no children marks it as good (leaf node)
• Returns 1, which is correct

5. Comparing with -1 Instead of Using a Flag

Using -1 as a sentinel value for first_child_subtree_size can be confusing and error-prone if subtree sizes could theoretically be negative (though they can't in this problem).

Alternative Approach:

def dfs(current_node, parent_node):
    child_sizes = []
    subtree_size = 1
  
    for child in adjacency_list[current_node]:
        if child != parent_node:
            child_size = dfs(child, current_node)
            child_sizes.append(child_size)
            subtree_size += child_size
  
    # Check if all sizes are equal (or no children)
    is_good = len(set(child_sizes)) <= 1
    self.good_nodes_count += is_good
  
    return subtree_size

This approach is cleaner but slightly less efficient due to set creation. The original approach with -1 is more efficient but requires careful handling.