"Count Nodes Equal to Average of Subtree" solution

Author: lancelote

src/count_nodes_equal_to_average_of_subtree.py

@@ -0,0 +1,26 @@
+from src.utils.binary_tree import TreeNode
+
+
+class Solution:
+    def averageOfSubtree(self, root: TreeNode | None) -> int:
+        count = 0
+
+        def dfs(node: TreeNode | None) -> tuple[int, int]:
+            nonlocal count
+
+            if node is None:
+                return 0, 0
+
+            left_count, left_total = dfs(node.left)
+            right_count, right_total = dfs(node.right)
+
+            local_count = 1 + left_count + right_count
+            local_total = node.val + left_total + right_total
+
+            if int(local_total / local_count) == node.val:
+                count += 1
+
+            return local_count, local_total
+
+        dfs(root)
+        return count

tests/test_count_nodes_equal_to_average_of_subtree.py

@@ -0,0 +1,16 @@
+import pytest
+
+from src.count_nodes_equal_to_average_of_subtree import Solution
+from src.utils.binary_tree import list_to_tree
+
+
+@pytest.mark.parametrize(
+    "in_list,expected",
+    (
+        ([4, 8, 5, 0, 1, None, 6], 5),
+        ([1], 1),
+    ),
+)
+def test_solution(in_list, expected):
+    root = list_to_tree(in_list)
+    assert Solution().averageOfSubtree(root) == expected