Program to count how many ways we can divide the tree into two trees in Python

Suppose we have a binary tree containing values 0, 1 and 2. The root has at least one 0 node and one 1 node. We need to find how many ways we can delete one edge such that the resulting two trees don't both contain 0 and 1 nodes.

In other words, after deleting an edge, each resulting subtree should contain either only 0s and 2s, or only 1s and 2s, but not both 0s and 1s together.

Example Tree Structure

The output will be 1 because we can only delete the edge between node 0 (root) and node 2, which separates 0s from 1s.

Algorithm Approach

We use a two-pass DFS approach ?

1. First pass: Count nodes (0s and 1s) in each subtree
2. Second pass: Check each edge to see if removing it creates valid separation

Implementation

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

class Solution:
    def solve(self, root):
        count = [0, 0, 0]  # Count of 0s, 1s, 2s in entire tree
        
        def dfs(node):
            if node:
                pre = count[:]  # Save current state
                dfs(node.left)
                dfs(node.right)
                count[node.val] += 1
                # Store count of 0s and 1s in this subtree
                node.count = [count[i] - pre[i] for i in range(2)]
        
        dfs(root)
        
        def dfs2(node, par=None):
            if node:
                if par is not None:
                    a0, a1 = node.count  # 0s and 1s in current subtree
                    b0, b1 = count[0] - a0, count[1] - a1  # 0s and 1s in remaining tree
                    # Check if removing edge creates valid separation
                    if (a0 == 0 or a1 == 0) and (b0 == 0 or b1 == 0):
                        self.ans += 1
                dfs2(node.left, node)
                dfs2(node.right, node)
        
        self.ans = 0
        dfs2(root)
        return self.ans

# Test the solution
ob = Solution()
root = TreeNode(0)
root.left = TreeNode(0)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)

print(ob.solve(root))

1

How It Works

The algorithm works in two phases ?

1. DFS1: Post-order traversal to count 0s and 1s in each subtree
2. DFS2: For each edge, check if removing it separates 0s from 1s

An edge is valid for removal if:

• One side has no 0s OR no 1s
• AND the other side has no 0s OR no 1s

Key Points

• Node value 2 can appear in either subtree after division
• We only need to separate nodes with values 0 and 1
• The solution uses bottom-up counting to efficiently track subtree compositions

Conclusion

This solution efficiently finds valid edge deletions using two DFS passes. The first pass counts node types in subtrees, while the second pass checks separation validity for each edge.