[리트코드] 2265 - Count Nodes Equal to Average of Subtree

2265. Count Nodes Equal to Average of Subtree

Medium

---

Given the root of a binary tree, return the number of nodes where the value of the node is equal to the average of the values in its subtree.

Note:

• The average of n elements is the sum of the n elements divided by n and rounded down to the nearest integer.
• A subtree of root is a tree consisting of root and all of its descendants.

 

Example 1:

Input: root = [4,8,5,0,1,null,6]
Output: 5
Explanation: 
For the node with value 4: The average of its subtree is (4 + 8 + 5 + 0 + 1 + 6) / 6 = 24 / 6 = 4.
For the node with value 5: The average of its subtree is (5 + 6) / 2 = 11 / 2 = 5.
For the node with value 0: The average of its subtree is 0 / 1 = 0.
For the node with value 1: The average of its subtree is 1 / 1 = 1.
For the node with value 6: The average of its subtree is 6 / 1 = 6.

Example 2:

Input: root = [1]
Output: 1
Explanation: For the node with value 1: The average of its subtree is 1 / 1 = 1.

 

Constraints:

• The number of nodes in the tree is in the range [1, 1000].
• 0 <= Node.val <= 1000

분류

트리, 깊이 우선 탐색, 이진 트리

---

문제 풀이

문제 분석

이 문제는 이진 트리의 루트가 주어졌을 때, 해당 노드의 값이 자신의 서브트리에 있는 모든 노드 값의 평균과 동일한 노드의 개수를 찾아 반환하도록 요구합니다. 여기서 서브트리의 평균은 서브트리에 있는 노드 값들의 총합을 노드의 개수로 나눈 후, 결과를 가장 가까운 정수로 **내림(rounded down)**한 값입니다. 즉, 나눗셈은 정수 나눗셈(//)을 사용해야 합니다.

입력은 이진 트리의 root 노드이며, 출력은 조건을 만족하는 노드의 총 개수입니다. 트리의 노드 개수는 1개부터 1000개까지, 노드 값은 0부터 1000까지의 범위를 가집니다.

접근 방법

각 노드에 대해 그 노드의 서브트리 전체의 합과 서브트리의 노드 개수를 알아야 평균을 계산할 수 있습니다. 이 정보는 자식 노드들로부터 얻어와야 하므로, 트리를 '아래에서 위로' 탐색하는 방식이 효과적입니다. 이러한 탐색 방식은 깊이 우선 탐색(DFS)의 일종인 **후위 순회(post-order traversal)**와 잘 맞습니다.

우리의 DFS 함수는 다음과 같은 정보를 반환해야 합니다:

1. 현재 노드를 루트로 하는 서브트리 내 모든 노드 값의 합.
2. 현재 노드를 루트로 하는 서브트리 내 모든 노드의 개수.

이 두 가지 정보를 재귀적으로 계산하면서, 각 노드마다 서브트리의 평균을 구하고 해당 노드의 값과 비교하여 조건을 만족하면 전역 카운터를 증가시키는 방식으로 문제를 해결할 수 있습니다.

구현 설명

1단계: 전역 카운터 초기화

• averageOfSubtree 함수가 호출될 때마다 global count 변수를 0으로 초기화합니다. 이는 여러 테스트 케이스를 실행할 때 이전 테스트 케이스의 결과가 다음 테스트 케이스에 영향을 미치지 않도록 하기 위함입니다. count는 조건을 만족하는 노드의 총 개수를 저장할 변수입니다.

2단계: DFS 헬퍼 함수 정의

• DFS(node)라는 재귀 헬퍼 함수를 정의합니다. 이 함수는 각 노드의 서브트리에 대한 합계와 노드 개수를 계산하고 반환합니다.
• 기저(Base) 조건: 만약 node가 None (즉, 빈 서브트리)이면, 합계와 노드 개수 모두 0이므로 (0, 0)을 반환합니다.
• 재귀(Recursive) 단계:
  • 먼저 왼쪽 자식 서브트리에 대해 DFS(node.left)를 호출하여 left_sum과 left_cnt를 얻습니다.
  • 다음으로 오른쪽 자식 서브트리에 대해 DFS(node.right)를 호출하여 right_sum과 right_cnt를 얻습니다.
  • 현재 노드를 포함한 전체 서브트리의 총합 total은 node.val + left_sum + right_sum이 됩니다.
  • 현재 노드를 포함한 전체 서브트리의 노드 개수 cnt는 1 + left_cnt + right_cnt가 됩니다 (여기서 1은 현재 노드 자신을 의미).

3단계: 조건 확인 및 카운터 증가

• 2단계에서 계산한 total과 cnt를 사용하여 현재 서브트리의 평균을 계산합니다. 문제에서 평균은 내림 처리되므로, Python의 정수 나눗셈 연산자 //를 사용하여 total // cnt를 수행합니다.
• 만약 이 계산된 평균이 현재 노드의 값 node.val과 같다면, global count 변수를 1 증가시킵니다.
• 마지막으로, 현재 노드의 서브트리에 대한 (total, cnt) 정보를 부모 노드로 반환합니다.

4단계: DFS 호출 및 결과 반환

• averageOfSubtree 함수 내에서 DFS(root)를 한 번 호출하여 전체 트리를 탐색하고 계산을 시작합니다.
• DFS(root) 호출이 완료되면, global count 변수에는 조건을 만족하는 모든 노드의 개수가 누적되어 있을 것입니다. 최종적으로 이 count 값을 반환합니다.

⏱복잡도 분석

• 시간 복잡도: O(N)
  • DFS 함수는 트리의 각 노드를 정확히 한 번 방문합니다. 각 노드에서 수행되는 작업(덧셈, 비교, 할당)은 상수 시간 O(1)이 걸립니다. 따라서 전체 시간 복잡도는 트리의 노드 수 N에 비례합니다.
• 공간 복잡도: O(H)
  • 재귀 호출 스택의 깊이가 트리의 높이 H에 비례합니다. 최악의 경우(예: 한쪽으로 치우친 트리, skewed tree) 트리의 높이 H가 노드 수 N과 같아질 수 있으므로 O(N)입니다. 최상의 경우(예: 균형 잡힌 트리, balanced tree) H는 log N이므로 O(log N)입니다.

핵심 포인트

1. 후위 순회(Post-order Traversal) DFS: 각 노드의 서브트리 합계와 노드 개수를 계산하기 위해 자식 노드들의 정보를 먼저 얻는 bottom-up 방식의 DFS가 필수적입니다.
2. 두 가지 값 반환: DFS 함수는 단순한 불리언 값이나 단일 값 대신, 서브트리의 총합과 노드 개수라는 두 가지 정보를 튜플 형태로 함께 반환해야 합니다.
3. 전역 카운터와 초기화: 조건을 만족하는 노드의 총 개수를 추적하기 위해 전역 변수(또는 클로저 내의 비지역 변수)를 사용하고, 각 테스트 케이스마다 정확한 결과를 얻기 위해 이를 초기화하는 것이 중요합니다.
4. 정수 나눗셈 //: 문제에서 평균을 "내림"하도록 요구하므로, 파이썬의 // 연산자를 사용하여 정수 나눗셈을 정확하게 수행해야 합니다.

풀이 코드

# 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
count = 0
class Solution:
    def averageOfSubtree(self, root: TreeNode) -> int:
        global count
        count = 0
        
        def DFS(node):
            global count
            if not node:
                return 0, 0

            left_sum, left_cnt = DFS(node.left)
            right_sum, right_cnt = DFS(node.right)

            total = node.val + left_sum + right_sum
            cnt = 1 + left_cnt + right_cnt

            if total // cnt == node.val:
                count += 1

            return total, cnt
        
        DFS(root)
        return count