algorithms

Algorithm implementations
git clone git://git.laack.co/algorithms.git
Log | Files | Refs | README

count-nodes-equal-to-average-of-subtree.py (896B)

      1 # Definition for a binary tree node.
      2 # class TreeNode:
      3 #     def __init__(self, val=0, left=None, right=None):
      4 #         self.val = val
      5 #         self.left = left
      6 #         self.right = right
      7 
      8 # idea
      9     # backtracking
     10     # average of this is  sum of values / count of subtrees
     11     # return this up the stack at each step
     12     # have shared state that tracks count then return that
     13 
     14 
     15 class Solution:
     16     def recurse(self, nd):
     17         if nd is None:
     18             return (0,0)
     19         left = self.recurse(nd.left)
     20         right = self.recurse(nd.right)
     21         
     22         sum_val = (left[1] + right[1]) + nd.val
     23         count = left[0] + right[0] + 1
     24         avg = sum_val // count
     25         if avg == nd.val:
     26             self.count += 1
     27         return (count, sum_val)
     28 
     29     def averageOfSubtree(self, root: TreeNode) -> int:
     30         self.count = 0
     31         self.recurse(root)
     32         return self.count