If we remove some edge from parent A to child B, then the subtree at B must have sum equal to the sum of all the nodes divided by 2.
Let's remember the subtree sum of every child node (stored in
sums) as we calculate the subtree sum of the root. After, we simply check that some subtree could equal half the total. We must also be careful that there is an edge to delete, hence the check for "
not root.left and not root.right".
def checkEqualTree(self, root): sums = set() def sum_(node, child = False): if not node: return 0 ans = node.val + sum_(node.left, True) + sum_(node.right, True) if child: sums.add(ans) return ans total = sum_(root) if total % 2 or not root.left and not root.right: return False return total // 2 in sums
@StefanPochmann Corrected it, thanks.