easyTreesPattern: DFS
Count Nodes in Each Subtree Solution
Problem Statement
Given the root of a tree, return a list of integers where the element at index `i` is the total number of nodes in the subtree rooted at node `i`. The tree is represented using an adjacency list where `adj[i]` contains a list of all nodes directly connected to node `i`. Assume nodes are labeled from 0 to n-1, where n is the total number of nodes. The input will be the number of nodes `n` and an adjacency list `adj`. You can assume node 0 is the root of the tree.
Examples
Example 1:
Input:n = 5, adj = [[1, 2], [0, 3, 4], [0], [1], [1]]
Output:[5, 3, 1, 1, 1]
Explanation: The tree structure is: 0 is connected to 1 and 2. 1 is connected to 0, 3, and 4. 2 is connected to 0. 3 is connected to 1. 4 is connected to 1. Assuming 0 is the root: Subtree rooted at 0 has nodes {0, 1, 2, 3, 4} (count=5). Subtree rooted at 1 has nodes {1, 3, 4} (count=3). Subtree rooted at 2 has nodes {2} (count=1). Subtree rooted at 3 has nodes {3} (count=1). Subtree rooted at 4 has nodes {4} (count=1).
Example 2:
Input:n = 3, adj = [[1], [0, 2], [1]]
Output:[3, 2, 1]
Explanation: The tree structure is: 0 is connected to 1. 1 is connected to 0 and 2. 2 is connected to 1. Assuming 0 is the root: Subtree rooted at 0 has nodes {0, 1, 2} (count=3). Subtree rooted at 1 has nodes {1, 2} (count=2). Subtree rooted at 2 has nodes {2} (count=1).
Example 3:
Input:n = 1, adj = [[]]
Output:[1]
Explanation: A single node tree. The subtree rooted at node 0 has node {0} (count=1).
Constraints
- 1 <= n <= 10^5
- The input `adj` will represent a valid tree structure.
- Each node will have a degree of at least 1, except for the root in a single-node tree.
- 0 <= adj[i].length <= n-1
- 0 <= adj[i][j] < n
Time: O(N + E) Space: O(N)
Use a single Depth First Search (DFS) starting from the root (node 0). During the post-order traversal phase of the DFS, for each node, sum the counts returned by its children's recursive calls and add 1 (for the node itself). Store this sum as the subtree count for the current node.
Run, Test & Submit Code
Ready to practice this challenge? Launch our interactive compilation environment with compiler validation.
Solve on Interactive WorkspaceTested Solutions
def count_nodes_in_subtree(n: int, adj: list[list[int]]) -> list[int]:
subtree_counts = [0] * n
visited = [False] * n
def dfs(node):
visited[node] = True
count = 1 # Count the current node
for neighbor in adj[node]:
if not visited[neighbor]:
count += dfs(neighbor)
subtree_counts[node] = count
return count
# Start DFS from the root (node 0)
dfs(0)
return subtree_counts