LeetCode 637. Average of Levels in Binary Tree
# 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 averageOfLevels(self, root: Optional[TreeNode]) -> List[float]:
ans = []
queue = []
queue.append([root, 1])
i = 0
sum1 = 0
count = 0
dep = 1
while i < len(queue):
node, depth = queue[i]
if dep != depth:
ans.append(sum1/count)
sum1 = 0
count = 0
dep += 1
sum1 += node.val
count += 1
if node.left:
queue.append([node.left, depth+1])
if node.right:
queue.append([node.right, depth+1])
i += 1
ans.append(sum1/count)
return ans
各levelごとの平均を求めたいのでbfsを使う
イメージはqueueだけど、いちいち先頭の要素を取り除くと計算時間がかかるので、一つの要素を見たら次の要素に移る方針で
depという今調べているlevelの深さを管理する変数を導入して、depとnodeのdepthが異なれば、depのlevelの平均を出す。
全てのノードに対して、そのノードのvalをsum1に足すことを忘れない!!
一番下のlevelはwhile文だけでは平均が出ないので(次のlevelに行ったときに平均を出すコードなので)、while文を抜けた後に付け足す!!