easyTreesPattern: DFS

Subtree Node Count Solution

Problem Statement

Given the root of a binary tree, return an array where each element at index `i` represents the total number of nodes in the subtree rooted at the `i`-th node (inclusive). The nodes are indexed level by level, from left to right. If the tree is empty, return an empty array.

Examples

Example 1:

Input:{"nodes":[1,2,3,4,5,6,7],"left":[1,3,5,-1,-1,-1,-1],"right":[2,4,6,-1,-1,7,-1]}

Output:[7,4,2,1,1,2,1]

Explanation: The tree nodes are: 1(root), 2(left of 1), 3(right of 1), 4(left of 2), 5(right of 2), 6(left of 3), 7(right of 3). Node 7 is the right child of node 3. Node 1's subtree has 7 nodes. Node 2's subtree has 4 nodes. Node 3's subtree has 2 nodes. Node 4's subtree has 1 node. Node 5's subtree has 1 node. Node 6's subtree has 2 nodes (itself and node 7). Node 7's subtree has 1 node.

Example 2:

Input:{"nodes":[10,20],"left":[20,-1],"right":[-1,-1]}

Output:[2,1]

Explanation: Node 10 is the root. Node 20 is its left child. Node 10's subtree has 2 nodes. Node 20's subtree has 1 node.

Example 3:

Input:{"nodes":[],"left":[],"right":[]}

Output:[]

Explanation: An empty tree has no nodes, so the output is an empty array.

Constraints

The number of nodes in the tree is between 0 and 1000.
The value of each node is unique.
The `left` and `right` arrays will have the same length as the `nodes` array.
A value of -1 in `left` or `right` indicates no child.

Time: O(N) Space: O(N)

Utilize a single Depth First Search (DFS) traversal. For each node, recursively calculate the size of its left and right subtrees. The size of the current node's subtree is 1 plus the sum of its children's subtree sizes. Store these counts in an array that maps node indices to their subtree sizes.

Run, Test & Submit Code

Ready to practice this challenge? Launch our interactive compilation environment with compiler validation.

Solve on Interactive Workspace

Tested Solutions

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def subtreeNodeCount(root: TreeNode) -> list[int]:
    if not root:
        return []

    # Helper function to build tree from example input format
    def build_tree(nodes_data, left_data, right_data):
        if not nodes_data:
            return None
        
        nodes = [None] * len(nodes_data)
        node_val_to_obj = {}
        for i, val in enumerate(nodes_data):
            if val != -1:
                nodes[i] = TreeNode(val)
                node_val_to_obj[val] = nodes[i]
        
        for i, node in enumerate(nodes):
            if node:
                # For simplicity, assuming indices in left_data and right_data correspond to indices in nodes_data
                # This implies the input format needs careful handling based on level order
                # The provided example structure needs to be mapped to actual tree connections.
                # The problem statement implies level-order indexing, which is crucial.
                # Let's assume the problem means node values map directly for simplicity in this correction.
                # A more robust tree builder would parse level by level.
                pass # The original code's build_tree logic was problematic for general cases
        
        # A proper tree builder based on level-order indexing would be more complex.
        # For the purpose of fixing the core logic, let's assume the root is passed and has correct structure.
        return nodes[0] if nodes else None

    # The core logic to count subtree nodes
    node_to_index = {}
    queue = [(root, 0)]
    index_counter = 0
    all_nodes_in_order = []
    
    while queue:
        current_node, current_index = queue.pop(0)
        if current_node:
            all_nodes_in_order.append(current_node.val)
            node_to_index[current_node] = index_counter
            index_counter += 1
            
            # Enqueue children for level-order traversal to get indices
            # This assumes the tree is already correctly built and we are mapping nodes to indices
            queue.append((current_node.left, -1)) # Dummy index for children, will be re-indexed
            queue.append((current_node.right, -1))
        else:
            # Handle None nodes if necessary for structure, but they don't contribute to count
            pass

    # Re-map to actual order of nodes encountered in a level-order traversal to get correct indices
    actual_nodes_in_level_order = []
    q_for_indices = [root]
    while q_for_indices:
        curr = q_for_indices.pop(0)
        if curr:
            actual_nodes_in_level_order.append(curr)
            q_for_indices.append(curr.left)
            q_for_indices.append(curr.right)

    n = len(actual_nodes_in_level_order)
    if n == 0: return []
    
    # Create a mapping from node object to its index in the level-order list
    node_obj_to_level_index = {node: i for i, node in enumerate(actual_nodes_in_level_order)}
    
    subtree_counts = [0] * n

    def dfs(node):
        if not node:
            return 0
        
        left_count = dfs(node.left)
        right_count = dfs(node.right)
        
        current_count = 1 + left_count + right_count
        
        # Get the index of the current node in the level-order traversal
        current_node_level_index = node_obj_to_level_index[node]
        subtree_counts[current_node_level_index] = current_count
        
        return current_count

    dfs(root)
    return subtree_counts