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文を抜けた後に付け足す!!