2265. Count Nodes Equal to Average of Subtree

[zeikar]

Problem Link

https://leetcode.com/problems/count-nodes-equal-to-average-of-subtree/

Problem Summary

이진 트리에서 자기 값이 서브트리 전체 값의 평균(내림)과 같은 노드의 개수를 구하는 문제.

Solution

노드마다 서브트리를 다시 돌면 O(n^2)이다.
평균을 구하려면 서브트리의 합과 개수만 있으면 되고, 둘 다 자식 값으로 바로 구할 수 있다. 즉, 후위 순회로 (합, 개수)를 올려주면 한 번만 돌아도 된다.

정답 개수도 같은 튜플에 넣어서 올려주면 따로 변수가 필요 없다.

search(node) = (서브트리 합, 서브트리 노드 수, 서브트리 안의 정답 개수)
빈 노드는 (0, 0, 0)

평균은 문제에서 내림이라고 했고 값이 전부 0 이상이라 // 쓰면 된다.

참고로 노드가 최대 1000개라 한쪽으로 쏠린 트리면 재귀 깊이도 1000인데, 파이썬 기본 재귀 한도도 1000이다 (실제로 깊이 996에서 터진다). 통과는 했으니 그런 입력은 없는 듯?

시간복잡도는 O(n)

Source 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:

        def search(node: TreeNode):
            if node == None:
                return (0, 0, 0)

            left = search(node.left)
            right = search(node.right)

            weights = left[0] + right[0] + node.val
            cnt = left[1] + right[1] + 1
            res = left[2] + right[2]

            if weights // cnt == node.val:
                res += 1

            return (weights, cnt, res)

        return search(root)[2]