[LeetCode] 637 | Average of Levels in Binary Tree | Python

---

문제 링크: https://leetcode.com/problems/average-of-levels-in-binary-tree/

---

문제

Given the root of a binary tree, return the average value of the nodes on each level in the form of an array. Answers within $10^{-5}$ of the actual answer will be accepted.

---

조건

• The number of nodes in the tree is in the range [1, $10^4$].
• $-2^{31}$ <= Node.val <= $2^{31}$ - 1

---

내 풀이

Tree구조는 recursive한 함수를 짜는 것이 좋다.

recursive하게 훑으며 values라는 리스트에 각 depth별 값을 모은다.

나중에 depth별로 평균을 구한 후 이를 리턴하면 된다.

---

코드

# 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]:
        values = []
        
        def search(root, depth):
            nonlocal values
            if not root:
                return
            
            if len(values) < depth+1:
                values.append([])
            
            values[depth].append(root.val)
            
            search(root.left, depth+1)
            search(root.right, depth+1)
            
        search(root, 0)
        ans = [sum(x)/len(x) for x in values]
        
        return ans