diff --git a/2347-count-nodes-equal-to-average-of-subtree/2347-count-nodes-equal-to-average-of-subtree.py b/2347-count-nodes-equal-to-average-of-subtree/2347-count-nodes-equal-to-average-of-subtree.py new file mode 100644 index 0000000..98c0183 --- /dev/null +++ b/2347-count-nodes-equal-to-average-of-subtree/2347-count-nodes-equal-to-average-of-subtree.py @@ -0,0 +1,23 @@ +class Solution: + def averageOfSubtree(self, root: TreeNode) -> int: + result = 0 + + def traverse(node): + nonlocal result + + if not node: + return 0, 0 + + left_sum, left_count = traverse(node.left) + right_sum, right_count = traverse(node.right) + + curr_sum = node.val + left_sum + right_sum + curr_count = 1 + left_count + right_count + + if curr_sum // curr_count == node.val: + result += 1 + + return curr_sum, curr_count + + traverse(root) + return result \ No newline at end of file