2265. Count Nodes Equal to Average of Subtree (Tree, Depth-First Search, Binary Tree)

링크 - https://leetcode.com/problems/count-nodes-equal-to-average-of-subtree/description/?envType=company&envId=facebook&favoriteSlug=facebook-thirty-days

난이도 - Medium

Intuition

이 문제는 dfs로 풀 수 있다. dfs로 현재 값들의 총합과, 총 노드의 개수를 반환하면 된다.

Code

# 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: TreeNode) -> int:
        ans = 0
        def dfs(node):
            if not node:
                return (0, 0)

            nonlocal ans

            left_sum, left_num_nodes = dfs(node.left)
            right_sum, right_num_nodes = dfs(node.right)

            total_val = node.val + left_sum + right_sum
            num_nodes = left_num_nodes + right_num_nodes + 1

            avg_val = total_val // num_nodes
            if node.val == avg_val:
                ans += 1

            return (total_val, num_nodes)

        dfs(root)
        return ans

Complexity

Time Complexity: O(n) n개의 노드를 방문한다.

Space Complexity: O(n) dfs의 recursive때문에 n개의 공간이 필요하다.