#3249

Medium

auto_awesomeBinary-tree traversal and state tracking

LeetCode Problem Workspace

Count the Number of Good Nodes

Determine how many nodes in a binary tree have all child subtrees of equal size using DFS traversal efficiently.

Problem Statement

You are given an undirected tree with n nodes labeled from 0 to n - 1, rooted at node 0, represented as a 2D integer array edges where edges[i] = [ai, bi] indicates an edge between nodes ai and bi. A node is considered good if all subtrees rooted at its children have the same size.

Implement a function that returns the total number of good nodes in the tree. Use DFS to traverse and track subtree sizes carefully to ensure nodes are correctly evaluated.

Examples

Example 1

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

Output: 7

All of the nodes of the given tree are good.

Example 2

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

Output: 6

There are 6 good nodes in the given tree. They are colored in the image above. 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]] Output: 12 Explanation: All nodes except node 9 are good.

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]]

Output: 12

All nodes except node 9 are good.

Constraints

2 <= n <= 105
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
The input is generated such that edges represents a valid tree.

Solution Approach

DFS Traversal with Subtree Size Tracking

Perform a recursive depth-first search from the root. For each node, calculate the sizes of its children subtrees. Compare these sizes to determine if the node is good. Maintain a global counter to sum all good nodes during traversal.

Edge List to Adjacency List Conversion

Convert the given edges array into an adjacency list to facilitate efficient tree traversal. Ensure that the traversal avoids revisiting parent nodes by keeping track of visited nodes or parent pointers.

Early Termination and Optimization

If a node's children subtrees already differ in size, mark it as not good immediately to avoid unnecessary computations. Combine subtree size calculation with good node checking in one DFS pass for O(n) time complexity.

Complexity Analysis

Metric	Value
Time	Depends on the final approach
Space	Depends on the final approach

Time complexity is O(n) because each node is visited once during DFS. Space complexity is O(n) for the adjacency list and recursion stack in the worst case, depending on the tree depth.

What Interviewers Usually Probe

Emphasizes DFS recursion and state tracking.
Expects handling of edge cases where nodes have varying numbers of children.
Checks for optimization in counting good nodes during traversal instead of post-processing.

Common Pitfalls or Variants

Common pitfalls

Failing to track parent nodes can lead to revisiting edges and incorrect counts.
Incorrectly calculating subtree sizes can misclassify good nodes.
Not handling leaf nodes separately may cause off-by-one errors in counting.

Follow-up variants

Count nodes where all subtrees differ in size instead of being equal.
Apply the good node definition to an N-ary tree instead of a binary tree.
Return a list of all good node indices rather than just the count.

FAQ

What defines a good node in Count the Number of Good Nodes?

A good node has all its children subtrees with exactly the same size. This requires calculating subtree sizes during DFS.

Can this problem be solved without recursion?

Yes, an iterative DFS using a stack is possible, but you must track subtree sizes and parent relationships carefully.

Why is DFS preferred for this problem?

DFS allows bottom-up computation of subtree sizes efficiently, which is essential to determine if a node is good.

What is the maximum input size for this tree problem?

The tree can have up to 10^5 nodes, so O(n) solutions with careful space management are required.

Are leaf nodes always considered good nodes?

Yes, leaf nodes have no children, so they trivially satisfy the condition for good nodes.

terminal

Solution

Solution 1: DFS

First, we construct the adjacency list $\textit{g}$ of the tree based on the given edges $\textit{edges}$, where $\textit{g}[a]$ represents all the neighboring nodes of node $a$.

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

class Solution:
    def countGoodNodes(self, edges: List[List[int]]) -> int:
        def dfs(a: int, fa: int) -> int:
            pre = -1
            cnt = ok = 1
            for b in g[a]:
                if b != fa:
                    cur = dfs(b, a)
                    cnt += cur
                    if pre < 0:
                        pre = cur
                    elif pre != cur:
                        ok = 0
            nonlocal ans
            ans += ok
            return cnt

        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        ans = 0
        dfs(0, -1)
        return ans

Continue Practicing

#3241 time taken to mark all nodes #3203 find minimum diameter after merging two trees #3319 k th largest perfect subtree size in binary tree