diff options
Diffstat (limited to 'cnt_node_equal_subtree.py')
-rw-r--r-- | cnt_node_equal_subtree.py | 42 |
1 files changed, 42 insertions, 0 deletions
diff --git a/cnt_node_equal_subtree.py b/cnt_node_equal_subtree.py new file mode 100644 index 0000000..70c6ab9 --- /dev/null +++ b/cnt_node_equal_subtree.py @@ -0,0 +1,42 @@ +# 2265. Count Nodes Equal to Average of Subtree + +""" +given the root of a binary tree, return the number of nodes where the value +of the node is equal to the average of the values in its subtree. note, the +average of 'n' elements is the sum of the n elements divided by n and rounded +down to the nearest integer. also, a subtree of root is a tree consisting of +root and all of its descendants +""" + + +class TreeNode(object): + def __init__(self, val=0, left=None, right=None): + self.val = val + self.left = left + self.right = right + + +class Solution(object): + def averageOfSubtree(self, root): + ans = 0 + + def traverse(node): + nonlocal ans + if not node: + return 0, 0 + left_sum, left_count = traverse(node.left) + right_sum, right_count = traverse(node.right) + s = node.val + left_sum + right_sum + c = 1 + left_count + right_count + if s // c == node.val: + ans += 1 + return s, c + + traverse(root) + return ans + + +if __name__ == "__main__": + obj = Solution() + print(obj.averageOfSubtree(root=[4, 8, 5, 0, 1, null, 6])) + print(obj.averageOfSubtree(root=[1])) |