LeetCode 3249 - Count the Number of Good Nodes

LeetCode Problem 3249

Difficulty: 🟡 Medium
Topics: Tree, Depth-First Search

Solution

Problem Understanding

We are given an undirected tree with n nodes labeled from 0 to n - 1. The tree is rooted at node 0, which means every node except the root has exactly one parent, and each node may have zero or more children.

The task is to count how many nodes are considered "good".

A node is called good if all of the subtrees rooted at its direct children have the same size.

To understand this precisely, consider a node u with children v1, v2, v3. For each child, we compute the size of that child's subtree, meaning the total number of nodes reachable below that child including the child itself. If all of those subtree sizes are equal, then node u is good.

Leaf nodes are always good because they have no children. Since there are no subtree sizes to compare, the condition is vacuously true.

The input is given as an edge list for an undirected tree:

• edges[i] = [a, b] means there is an undirected edge between nodes a and b

• Since the graph is guaranteed to be a valid tree:

• it contains exactly n - 1 edges

• it is connected

• it has no cycles

The constraints are important:

• n can be as large as 100000
• This immediately rules out repeated subtree recomputation
• Any solution worse than linear or near-linear time will likely time out

The main challenge is efficiently computing subtree sizes while simultaneously checking whether all child subtrees of each node are equal.

Several edge cases are important:

• A single chain of nodes, where every node has at most one child
• A star-shaped tree, where the root has many leaf children
• Nodes with exactly one child, which are always good because there is only one subtree size
• Deep trees that could cause recursion depth problems in some languages
• Highly unbalanced trees where subtree sizes vary dramatically

The problem guarantees the graph is a valid tree, so we do not need to handle disconnected graphs or cycles.

Approaches

Brute Force Approach

A straightforward approach is to evaluate every node independently.

For each node:

1. Identify all of its children
2. Compute the subtree size of each child
3. Compare the sizes
4. Count the node if all sizes match

The expensive part is subtree size computation. If we recompute subtree sizes from scratch for every node, the same portions of the tree get traversed many times.

For example, in a chain-like tree:

• Computing the subtree size of the root visits all n nodes
• Computing the subtree size of the next node visits n - 1 nodes
• Then n - 2, and so on

This leads to quadratic complexity in the worst case.

Although this method is logically correct, it is far too slow for n = 100000.

Optimal Approach

The key observation is that subtree sizes should only be computed once.

A depth-first search naturally supports this because subtree sizes depend on children, so we can process the tree bottom-up.

For each node during DFS:

1. Recursively compute subtree sizes of all children
2. Store those subtree sizes
3. Check whether all child subtree sizes are equal
4. Return the current node's subtree size to the parent

This works efficiently because:

• Every edge is traversed exactly twice
• Every subtree size is computed exactly once
• The equality check for a node only depends on its direct children

The algorithm becomes a classic postorder DFS computation.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n²)	O(n)	Recomputes subtree sizes repeatedly
Optimal	O(n)	O(n)	Single DFS computes all subtree sizes once

Algorithm Walkthrough

1. Build an adjacency list from the edge list.

Since the tree is undirected, every edge [u, v] is added in both directions. The adjacency list allows efficient traversal of neighbors. 2. Start a DFS from the root node 0.

During DFS, we pass the parent node so we do not traverse backward and create an infinite loop. 3. For the current node, initialize:

• subtree_size = 1, counting the node itself
• a variable to track child subtree sizes
• a flag indicating whether all child subtree sizes match

4. Traverse all neighbors.

For each neighbor:

• skip the parent
• recursively compute the child's subtree size
• add that size to the current node's subtree size

5. Compare child subtree sizes.

The first child establishes the expected subtree size.

Every later child must match it.

If any child subtree size differs, the current node is not good. 6. After processing all children:

• if all child subtree sizes matched, increment the answer
• return the current node's subtree size to the parent

7. Continue until DFS finishes.

Since DFS processes children before parents, every subtree size is already known when needed.

Why it works

The algorithm relies on postorder DFS traversal.

When processing a node, all child subtree sizes have already been computed correctly by recursive calls. Because the subtree size of a node depends only on its descendants, this guarantees correctness.

A node is counted as good exactly when all child subtree sizes are equal. Since every child subtree size is computed accurately and compared exactly once, the final count is correct.

Python Solution

from typing import List
from collections import defaultdict

class Solution:
    def countGoodNodes(self, edges: List[List[int]]) -> int:
        n = len(edges) + 1

        graph = defaultdict(list)

        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        good_nodes = 0

        def dfs(node: int, parent: int) -> int:
            nonlocal good_nodes

            subtree_size = 1
            expected_size = -1
            is_good = True

            for neighbor in graph[node]:
                if neighbor == parent:
                    continue

                child_size = dfs(neighbor, node)

                if expected_size == -1:
                    expected_size = child_size
                elif expected_size != child_size:
                    is_good = False

                subtree_size += child_size

            if is_good:
                good_nodes += 1

            return subtree_size

        dfs(0, -1)

        return good_nodes

The implementation begins by constructing an adjacency list because the input tree is undirected. This structure enables efficient traversal from any node to its neighbors.

The DFS function returns the size of the subtree rooted at the current node. This return value is essential because parents need child subtree sizes to determine whether they are good.

Inside DFS:

• subtree_size starts at 1 because every subtree includes the current node itself
• expected_size stores the first child subtree size encountered
• is_good tracks whether all child subtree sizes remain equal

For every child:

• recursively compute the child's subtree size
• compare it against the expected size
• accumulate it into the current subtree size

After all children are processed, the node is counted if all subtree sizes matched.

Finally, DFS starts from the root node 0, and the total number of good nodes is returned.

Go Solution

package main

func countGoodNodes(edges [][]int) int {
	n := len(edges) + 1

	graph := make([][]int, n)

	for _, edge := range edges {
		u := edge[0]
		v := edge[1]

		graph[u] = append(graph[u], v)
		graph[v] = append(graph[v], u)
	}

	goodNodes := 0

	var dfs func(int, int) int

	dfs = func(node int, parent int) int {
		subtreeSize := 1
		expectedSize := -1
		isGood := true

		for _, neighbor := range graph[node] {
			if neighbor == parent {
				continue
			}

			childSize := dfs(neighbor, node)

			if expectedSize == -1 {
				expectedSize = childSize
			} else if expectedSize != childSize {
				isGood = false
			}

			subtreeSize += childSize
		}

		if isGood {
			goodNodes++
		}

		return subtreeSize
	}

	dfs(0, -1)

	return goodNodes
}

The Go implementation follows the same logic as the Python version.

Some language-specific details are worth noting:

• Go uses slices for the adjacency list
• The DFS function is declared as a closure so it can recursively reference itself
• Integers in Go are sufficient because subtree sizes never exceed 100000
• Unlike Python, there is no need for nonlocal, since closures can directly modify outer variables like goodNodes

Worked Examples

Example 1

Input:

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

Tree structure:

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

DFS processes bottom-up.

Node	Child Subtree Sizes	Good?	Subtree Size
3	[]	Yes	1
4	[]	Yes	1
1	[1, 1]	Yes	3
5	[]	Yes	1
6	[]	Yes	1
2	[1, 1]	Yes	3
0	[3, 3]	Yes	7

Total good nodes = 7.

Example 2

Input:

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

Tree structure:

0
├── 1
│   ├── 2
│   │   ├── 3
│   │   │   ├── 4
│   │   │   └── 8
│   │   └── 7
│   └── 6
└── 5

Process nodes bottom-up.

Node	Child Subtree Sizes	Good?	Subtree Size
4	[]	Yes	1
8	[]	Yes	1
3	[1, 1]	Yes	3
7	[]	Yes	1
2	[3, 1]	No	5
6	[]	Yes	1
1	[5, 1]	No	7
5	[]	Yes	1
0	[7, 1]	No	9

Total good nodes = 6.

Example 3

Input:

edges = [[0,1],[1,2],[1,3],[1,4],[0,5],[5,6],[6,7],[7,8],[0,9],[9,10],[9,12],[10,11]]

Key subtree sizes:

Node	Child Subtree Sizes	Good?
1	[1, 1, 1]	Yes
5	[3]	Yes
6	[2]	Yes
7	[1]	Yes
9	[2, 1]	No
0	[4, 4, 4]	Yes

Every node except 9 is good.

Total good nodes = 12.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Every node and edge is processed once
Space	O(n)	Adjacency list and recursion stack

The DFS visits every node exactly once. Since a tree with n nodes has n - 1 edges, the total traversal work is linear.

The adjacency list stores all edges, requiring O(n) space. The recursion stack may also reach O(n) depth in a highly skewed tree.

Test Cases

from typing import List

class Solution:
    def countGoodNodes(self, edges: List[List[int]]) -> int:
        from collections import defaultdict

        n = len(edges) + 1

        graph = defaultdict(list)

        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        answer = 0

        def dfs(node: int, parent: int) -> int:
            nonlocal answer

            subtree_size = 1
            expected = -1
            good = True

            for neighbor in graph[node]:
                if neighbor == parent:
                    continue

                child_size = dfs(neighbor, node)

                if expected == -1:
                    expected = child_size
                elif expected != child_size:
                    good = False

                subtree_size += child_size

            if good:
                answer += 1

            return subtree_size

        dfs(0, -1)

        return answer

solver = Solution()

assert solver.countGoodNodes(
    [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]
) == 7  # perfectly balanced binary tree

assert solver.countGoodNodes(
    [[0,1],[1,2],[2,3],[3,4],[0,5],[1,6],[2,7],[3,8]]
) == 6  # unbalanced subtree sizes

assert solver.countGoodNodes(
    [[0,1],[1,2],[1,3],[1,4],[0,5],[5,6],[6,7],[7,8],[0,9],[9,10],[9,12],[10,11]]
) == 12  # only one bad node

assert solver.countGoodNodes(
    [[0,1]]
) == 2  # smallest valid tree

assert solver.countGoodNodes(
    [[0,1],[1,2],[2,3],[3,4]]
) == 5  # chain tree, every node good

assert solver.countGoodNodes(
    [[0,1],[0,2],[0,3],[0,4]]
) == 5  # star tree, root children all size 1

assert solver.countGoodNodes(
    [[0,1],[1,2],[1,3]]
) == 4  # balanced small tree

assert solver.countGoodNodes(
    [[0,1],[1,2],[2,3],[2,4]]
) == 5  # mixed branching

assert solver.countGoodNodes(
    [[0,1],[0,2],[2,3]]
) == 3  # root has unequal child subtree sizes

Test	Why
Perfect balanced binary tree	Verifies all nodes are counted
Unbalanced subtree example	Ensures unequal subtree sizes are detected
Only one bad node	Validates selective exclusion
Smallest valid tree	Tests minimum constraint
Chain tree	Confirms single-child nodes are always good
Star tree	Verifies equal leaf subtree sizes
Small balanced tree	Tests simple multi-child equality
Mixed branching tree	Validates recursive subtree accumulation
Unequal root subtrees	Ensures root comparison works correctly

Edge Cases

Leaf Nodes

Leaf nodes have no children, which means there are no subtree sizes to compare. A common mistake is to incorrectly treat them as invalid because there are zero child subtrees.

This implementation handles them naturally. The DFS loop never runs, is_good remains True, and the node is counted correctly.

Nodes With Exactly One Child

A node with one child is always good because there is only one subtree size. There is nothing to compare against.

Some implementations mistakenly require at least two matching subtree sizes. Here, the first child simply sets expected_size, and no mismatch can occur.

Highly Skewed Trees

A chain-shaped tree can create recursion depth concerns and inefficient repeated traversals in brute-force solutions.

This algorithm remains linear because each subtree size is computed once. Even in the worst-case chain structure, total work is still O(n).

Root Node Handling

The root has no parent, so DFS must avoid incorrectly revisiting nodes through undirected edges.

Passing the parent node during recursion prevents traversal back upward and guarantees correct tree traversal without cycles.