題目出處
2265. Count Nodes Equal to Average of Subtree
難度
medium
個人範例程式碼
# 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:
if not root:
return 0
self.ans = 0
self.dfs(root)
return self.ans
def dfs(self, node):
# end of recursion
if not node:
return 0, 0
# define and split
left_sum, left_cnt = self.dfs(node.left)
right_sum, right_cnt = self.dfs(node.right)
if (node.val + left_sum + right_sum) // (1 + left_cnt + right_cnt) == node.val:
self.ans += 1
return (node.val + left_sum + right_sum), (1 + left_cnt + right_cnt)
最近在練習程式碼本身就可以自解釋的 Coding style，可以嘗試直接閱讀程式碼理解
算法說明
處理 subtree 的內容，算出所有 subtree 的平均，
我們透過 dfs 往底部搜尋，回傳 sum 與 len
input handling
如果 no root ，return 0
Boundary conditions
dfs 結束條件，if not node，回傳 0, 0