[Leetcode/Python] 222. Count Complete Tree Nodes

문제 링크

Count Complete Tree Nodes - LeetCode

 

문제 설명

 완전 이진 트리의 root가 주어지면 트리의 노드 수를 반환합니다.

 Wikipedia에 따르면 마지막 레벨을 제외한 모든 레벨은 완전한 이진 트리로 완전히 채워져 있으며 마지막 레벨의 모든 노드는 가능한 한 가장 왼쪽에 있습니다. 마지막 레벨 h에서 포함하여 1~2^h 노드를 가질 수 있습니다.

 O(n) 시간 복잡도 미만으로 실행되는 알고리즘을 설계합니다.

 

입출력 예시

root	output
[1,2,3,4,5,6]	6
[]	0
[1]	1

 

제한 사항

• 트리의 노드 수는 [0, 5 * 10^4] 범위입니다.
• 0 <= Node.val <= 5 * 10^4
• 트리는 완전성을 보장합니다. (완전 이진 트리니 당연한 걸지도)

 

문제 풀이

 처음에는 간단히 DFS 알고리즘을 이용하려 했다. 그러나 해당문제는 시간복잡도가 O(n) 이하여야 한다. 간단히 말해서 O(log n) 수준으로 문제를 풀라는 뜻이다. 즉 모든 노드를 탐색하면 절대 안된다.

 

 여기서 나는 스택을 이용했다. 처음에는 (root, 1)을 스택에 넣었다. (타겟노드, n번째 수) 라는 뜻이다. 이후 스택에서 하나씩 꺼내서 (node.left, 2 * n), (node.right, 2 * n + 1)을 순서대로 넣었다. 왜 2 * n, 2 * n + 1이 되는지는 간단히 그림을 그려보면 알것이다.

 

 간단히 에시를 들어보겠다. 첫 테스트 케이스인 [1,2,3,4,5,6]을 넣어보자. 그럼 스택의 변화 과정은 다음과 같다.

1. [(root, 1)]
2. [(root.left, 2), (root.right, 3)]
3. [(root.left, 2), (root.right.left, 6)]

 이후 마지막에 6을 꺼내면 해당 노드는 left와 right가 존재하지 않는다. 그럼 마지막 수를 찾았다는 뜻이다. 이 경우 count에 6을 넣고 스택을 다시 하나씩 꺼내는 행위를 반복한다.

 

 만약 스택을 꺼내면서 count 보다 높은 수가 나왔다면, 그 수를 리턴해버리면 된다. 처음에 나온 count는 가장 우측의 수를 의미했지만, 완전 이진 트리이기 때문에 새로 나온 수는 가장 밑단 우측의 수이기 때문이다. 

 

 스택이 다 없어질때까지 count가 변하지 않는다면 이는 가장 밑단까지 꽉찬 완전 이진트리다. 내 알고리즘의 경우 이러한 경우 O(n)의 성능이 나오는게 단점이다.

 

소스코드

# 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 countNodes(self, root: Optional[TreeNode]) -> int:
        if not root:
            return 0
        stack = [(root, 1)]
        count = 0
        while stack:
            target = stack.pop()
            if not target[0].left:
                if count == 0:
                    count = target[1]
                elif count < target[1]:
                    count = target[1]
                    break
            else:
                stack.append((target[0].left, target[1] * 2))
                if target[0].right:
                    stack.append((target[0].right, target[1] * 2 + 1))
        return count

 

 

코딩 테스트, 어떻게 대비해야 할까?