Suppose we have a binary tree containing values 0, 1 and 2. The root has at least one 0 node and one 1 node. Now suppose there is an operation where we delete an edge in the tree and the tree becomes two different trees. We have to find the number of ways we can delete one edge such that none of the two trees contain both a 0 node and a 1 node.
So, if the input is like
then the output will be 1 as we can only delete the 0 to 2 edge.
To solve this, we will follow these steps −
Let us see the following implementation to get better understanding −
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] def dfs(node): if node: pre = count[:] dfs(node.left) dfs(node.right) count[node.val] += 1 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 b0, b1 = count - a0, count - a1 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 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))
root = TreeNode(0) root.left = TreeNode(0) root.right = TreeNode(2) root.right.left = TreeNode(1) root.right.right = TreeNode(1)