Teaching Kids Programming – Count Nodes Equal to Average of Subtree via Recursive Depth First Search Algorithm

Teaching Kids Programming: Videos on Data Structures and Algorithms

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.
A subtree of root is a tree consisting of root and all of its descendants.

Example 1:

Input: root = [4,8,5,0,1,null,6]
Output: 5
Explanation:
For the node with value 4: The average of its subtree is (4 + 8 + 5 + 0 + 1 + 6) / 6 = 24 / 6 = 4.
For the node with value 5: The average of its subtree is (5 + 6) / 2 = 11 / 2 = 5.
For the node with value 0: The average of its subtree is 0 / 1 = 0.
For the node with value 1: The average of its subtree is 1 / 1 = 1.
For the node with value 6: The average of its subtree is 6 / 1 = 6.

Example 2:
Input: root = [1]
Output: 1
Explanation: For the node with value 1: The average of its subtree is 1 / 1 = 1.

Constraints:
The number of nodes in the tree is in the range [1, 1000].
0 <= Node.val <= 1000

Hints:
What information do we need to calculate the average? We need the sum of the values and the number of values.
Create a recursive function that returns the size of a node’s subtree, and the sum of the values of its subtree.

Recursive Depth First Search Algorithm to Count Nodes Equal to Average of Subtree

In order to compute the average, we would need to know the sum of the nodes, and the number of nodes. We can return these information in a tuple by Recursive Depth First Search Algorithm. This is pretty much the same as Teaching Kids Programming – Recursive Depth First Search Algorithm to Compute the Max Average of a Binary SubTree

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def averageOfSubtree(self, root: Optional[TreeNode]) -> int:
        ans = 0
        
        def dfs(root):
            if not root:
                return 0, 0
            ls, ln = dfs(root.left)
            rs, rn = dfs(root.right)
            s = ls + rs + root.val
            n = ln + rn + 1
            if root.val == s // n:
                nonlocal ans
                ans += 1
            return s, n
        
        dfs(root)
        return ans

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def averageOfSubtree(self, root: Optional[TreeNode]) -> int:
        ans = 0
        
        def dfs(root):
            if not root:
                return 0, 0
            ls, ln = dfs(root.left)
            rs, rn = dfs(root.right)
            s = ls + rs + root.val
            n = ln + rn + 1
            if root.val == s // n:
                nonlocal ans
                ans += 1
            return s, n
        
        dfs(root)
        return ans

Each node is visisted at most once in Recursive DFS, so the time complexity is O(N) where N is the number of the nodes in the binary tree. The space complexity is O(N) given that the maximum calling depth is N.

With this approach aka Recursive DFS algorithm, we can solve this similar problem: Teaching Kids Programming – Count Nodes Equal to Sum of Descendants (Recursive Depth First Search Algorithm)

–EOF (The Ultimate Computing & Technology Blog) —