Problem Description

You are given a tree with nodes number of nodes, where the root is at node 0. The tree is represented using three pieces of information:

• The total number of nodes is nodes
• Each node i has a value given by value[i]
• Each node i (except the root) has a parent given by parent[i]

Your task is to remove every subtree whose sum of node values equals zero. After removing all such subtrees, return the count of nodes that remain in the tree.

A subtree includes a node and all of its descendants. When a subtree's total sum equals zero, the entire subtree (including the root of that subtree and all its descendants) should be removed from the tree.

For example, if a node and all its descendants sum to zero, that entire branch is deleted. The removal process should be applied recursively - if removing some subtrees causes a parent subtree to now sum to zero, that parent subtree should also be removed.

The solution uses DFS to traverse the tree from bottom to top, calculating the sum and node count for each subtree. When a subtree's sum equals zero, its node count is set to zero, effectively removing it from the final count.

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 involves a tree structure where nodes are connected through parent-child relationships. A tree is a special type of graph (specifically, a connected acyclic graph).

Is it a tree?

• Yes: The problem explicitly states we have a tree rooted at node 0, with parent-child relationships defined by the parent array.

DFS

• Yes: Since we identified this as a tree problem, the flowchart leads us directly to DFS as the solution approach.

Conclusion: The flowchart suggests using DFS (Depth-First Search) for this tree problem.

This makes perfect sense for Delete Tree Nodes because:

1. We need to traverse the entire tree to calculate subtree sums
2. We need to process children before parents (post-order traversal) to determine which subtrees sum to zero
3. DFS naturally handles the recursive nature of checking each subtree and propagating results upward
4. The problem requires exploring all paths from leaves to root to correctly identify and remove zero-sum subtrees

The DFS approach allows us to efficiently calculate subtree sums bottom-up and mark subtrees for removal when their sum equals zero.

Intuition

The key insight is that we need to know the sum of each subtree before deciding whether to keep or remove it. This naturally suggests a bottom-up approach where we process children before their parents.

Think about it this way: if we start from the root and go downward, we won't know whether to keep a node until we've explored all its descendants. But if we start from the leaves and work our way up, we can calculate each subtree's sum and immediately know if it should be removed.

When we visit a node during our traversal, we need to track two pieces of information:

1. The sum of all values in the subtree rooted at this node
2. The count of nodes that remain in this subtree after removal

Here's the clever part: when a subtree sums to zero, we don't actually delete it from the tree structure. Instead, we simply mark it as having 0 remaining nodes. This way, when the parent node asks "how many nodes are in my child's subtree?", it gets 0, effectively ignoring the removed subtree.

The process works like this:

• Start at a leaf node: its sum is just its own value, and it has 1 node
• Move to a parent: accumulate the sums and node counts from all children
• If the total sum equals zero, set the node count to 0 (marking the entire subtree for removal)
• Continue upward until we reach the root

By the time we finish processing the root, we have the total count of nodes that remain in the tree after all zero-sum subtrees have been "removed."

This approach elegantly handles cascading removals - if removing some child subtrees causes a parent subtree to sum to zero, that parent subtree will also be marked with 0 nodes, effectively removing it as well.

Learn more about Tree, Depth-First Search and Breadth-First Search patterns.

Solution Approach

The implementation follows the DFS pattern we identified, with a few key components:

Building the Tree Structure: First, we convert the parent array into an adjacency list representation using a defaultdict(list). For each node i (starting from 1, since 0 is the root with no parent), we add it to its parent's children list: g[parent[i]].append(i). This gives us easy access to all children of any node.

The DFS Function: We define a recursive function dfs(i) that processes node i and returns a tuple (sum, count) where:

• sum is the total sum of values in the subtree rooted at node i
• count is the number of nodes remaining in this subtree after removing zero-sum subtrees

Processing Each Node: For each node i, we:

1. Initialize s = value[i] (the sum starts with the node's own value)
2. Initialize m = 1 (count starts at 1 for the node itself)
3. Recursively process all children in g[i]:
   • Call dfs(j) for each child j
   • Add the child's subtree sum to our running total: s += t
   • Add the child's remaining node count to our count: m += n

Handling Zero-Sum Subtrees: After processing all children, we check if the total sum s equals 0. If it does, we set m = 0, effectively marking this entire subtree as removed. This is the crucial step - we don't actually delete nodes, we just report that this subtree has 0 remaining nodes.

Getting the Final Answer: We call dfs(0) to process the entire tree starting from the root. The function returns a tuple, and we extract the second element [1] which gives us the total count of remaining nodes.

The beauty of this approach is its simplicity - by returning both the sum and count in each recursive call, parent nodes can easily aggregate information from their children and make decisions about removal. The post-order nature of the traversal ensures that all descendants are processed before making decisions about ancestors, handling cascading removals naturally.

Example Walkthrough

Let's walk through a concrete example to see how the solution works.

Consider a tree with 5 nodes:

• Node values: [1, -1, 3, -3, 2]
• Parent array: [-1, 0, 0, 2, 2]

This creates the following tree structure:

       0(1)
      /    \
    1(-1)  2(3)
           /  \
         3(-3) 4(2)

Step 1: Build the adjacency list From the parent array, we create:

• g[0] = [1, 2] (node 0 has children 1 and 2)
• g[2] = [3, 4] (node 2 has children 3 and 4)

Step 2: Start DFS from root (node 0)

The DFS processes nodes in post-order (children before parents):

Processing node 1 (leaf):

• No children to process
• Sum = -1 (its own value)
• Count = 1 (just itself)
• Sum ≠ 0, so keep the count as 1
• Returns: (-1, 1)

Processing node 3 (leaf):

• No children to process
• Sum = -3
• Count = 1
• Sum ≠ 0, so keep the count as 1
• Returns: (-3, 1)

Processing node 4 (leaf):

• No children to process
• Sum = 2
• Count = 1
• Sum ≠ 0, so keep the count as 1
• Returns: (2, 1)

Processing node 2:

• Start with sum = 3 (its value), count = 1 (itself)
• Process child 3: adds (-3, 1) → sum becomes 3 + (-3) = 0, count becomes 1 + 1 = 2
• Process child 4: adds (2, 1) → sum becomes 0 + 2 = 2, count becomes 2 + 1 = 3
• Final sum = 2, count = 3
• Sum ≠ 0, so keep the count as 3
• Returns: (2, 3)

Processing node 0 (root):

• Start with sum = 1 (its value), count = 1 (itself)
• Process child 1: adds (-1, 1) → sum becomes 1 + (-1) = 0, count becomes 1 + 1 = 2
• Process child 2: adds (2, 3) → sum becomes 0 + 2 = 2, count becomes 2 + 3 = 5
• Final sum = 2, count = 5
• Sum ≠ 0, so keep the count as 5
• Returns: (2, 5)

Result: The function returns 5, meaning all 5 nodes remain in the tree.

Now let's see what happens if we modify node 4's value to be -2 instead of 2:

Modified tree:

       0(1)
      /    \
    1(-1)  2(3)
           /  \
         3(-3) 4(-2)

When we process node 2 now:

• Start with sum = 3, count = 1
• Add child 3: sum = 3 + (-3) = 0, count = 1 + 1 = 2
• Add child 4: sum = 0 + (-2) = -2, count = 2 + 1 = 3
• Final sum = -2, count = 3
• Returns: (-2, 3)

When we process node 0:

• Start with sum = 1, count = 1
• Add child 1: sum = 1 + (-1) = 0, count = 1 + 1 = 2
• Add child 2: sum = 0 + (-2) = -2, count = 2 + 3 = 5
• Final sum = -2, count = 5
• Returns: (-2, 5)

But if node 2's subtree had summed to 0:

• When processing node 2, if the sum equaled 0, we'd set count = 0
• This would return (0, 0) to node 0
• Node 0 would then only count itself and node 1, giving a final count of 2

This demonstrates how the algorithm elegantly handles subtree removal by simply setting the count to 0 when a subtree sums to zero, allowing parent nodes to naturally exclude removed subtrees from their counts.

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 DFS traversal. For each node, the operations performed are:

• Accessing the node's value: O(1)
• Iterating through its children: O(number of children)
• Performing constant-time arithmetic operations: O(1)

Since the total number of edges in a tree is n-1 (where n is the number of nodes), and each edge is traversed once when iterating through children across all nodes, the total time complexity is O(n).

Space Complexity: O(n)

The space complexity consists of:

• The adjacency list g (defaultdict): This stores all parent-child relationships. In the worst case, it stores n-1 edges, requiring O(n) space.
• The recursion call stack: In the worst case (a skewed tree), the maximum depth of recursion can be n, requiring O(n) space for the call stack.
• Local variables in each recursive call (s, m, t, n): These are constant space per call, contributing O(n) in total across all recursive calls.

Therefore, the overall space complexity is O(n).

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

Common Pitfalls

1. Incorrect Order of Processing - Not Using Post-Order Traversal

A common mistake is trying to process nodes in the wrong order, such as using pre-order traversal or attempting to delete nodes while traversing. This can lead to incorrect results because parent nodes need information from all their descendants before deciding whether to delete themselves.

Wrong Approach:

def dfs(node_idx):
    # WRONG: Checking deletion before processing children
    if value[node_idx] == 0:
        return (0, 0)
  
    subtree_sum = value[node_idx]
    node_count = 1
  
    for child_idx in adjacency_list[node_idx]:
        child_sum, child_count = dfs(child_idx)
        subtree_sum += child_sum
        node_count += child_count
  
    return (subtree_sum, node_count)

Why it fails: This approach only checks if the current node's value is 0, not the entire subtree sum. It would miss cases where a subtree sums to 0 due to positive and negative values canceling out.

2. Not Propagating Zero-Count Correctly

Another pitfall is forgetting that when a child subtree is deleted (count = 0), its sum should still be included in the parent's calculation, but its count should not affect whether the parent gets deleted.

Wrong Approach:

def dfs(node_idx):
    subtree_sum = value[node_idx]
    node_count = 1
  
    for child_idx in adjacency_list[node_idx]:
        child_sum, child_count = dfs(child_idx)
        # WRONG: Not adding child_sum if child was deleted
        if child_count > 0:
            subtree_sum += child_sum
            node_count += child_count
  
    if subtree_sum == 0:
        node_count = 0
      
    return (subtree_sum, node_count)

Why it fails: Even if a child subtree is deleted (count = 0), its sum value should still contribute to the parent's sum calculation. The parent needs the complete sum of all its descendants to determine if it should also be deleted.

3. Modifying Tree Structure During Traversal

Attempting to actually remove nodes from the adjacency list during traversal can cause issues:

Wrong Approach:

def dfs(node_idx):
    subtree_sum = value[node_idx]
    node_count = 1
  
    children_to_remove = []
    for child_idx in adjacency_list[node_idx]:
        child_sum, child_count = dfs(child_idx)
        if child_count == 0:
            # WRONG: Trying to modify structure during traversal
            children_to_remove.append(child_idx)
        subtree_sum += child_sum
        node_count += child_count
  
    # WRONG: Modifying the adjacency list
    for child in children_to_remove:
        adjacency_list[node_idx].remove(child)
  
    if subtree_sum == 0:
        node_count = 0
      
    return (subtree_sum, node_count)

Why it fails: There's no need to actually modify the tree structure. The algorithm works by tracking counts, not by physically removing nodes. Modifying the structure adds unnecessary complexity and doesn't improve the solution.

4. Edge Case: Single Node with Zero Value

Not handling the case where the root itself has value 0 and no children:

Correct Solution Handles This: The provided solution correctly handles this case because it checks the sum after processing all children (even if there are none), so a root with value 0 will correctly return count = 0.