[Leetcode] 110. Balanced Binary Tree

문제

• 관련 토픽: DFS, Binary Tree, Tree
• 난이도: Easy
• 링크: https://leetcode.com/problems/balanced-binary-tree/description/

요구사항

• 주어진 Binary Tree가 height balanced 인 상태인지 확인
  • height balanced = 두 서브트리의 높이의 차이가 1이상이 아닌 트리

조건

• 노드 개수 범위 [0, 5000]
• 10^4 <= node.val <= 10^4 -> O(N) 시간복잡도 이전으로 풀어야 한다

풀이

• 모든 노드의 왼쪽 서브트리와 오른쪽 서브트리의 높이 차이를 확인
• 높이 차이가 1을 초과하면, 트리는 불균형하다고 판단, return False
• DFS방식으로 각 노드의 서브트리 높이를 계산하며 균형여부를 확인
• Stack으로 DFS를 풀이하니, 각 노드에 방문했는지 여부를 저장하는 visited 배열 필요
• 현재 노드의 왼쪽 서브트리, 오른쪽 서브트리의 높이를 확인해야하기 때문에 (재귀의 경우는 반환값으로), heights라는 배열 필요
• 방문한 곳의 왼쪽 서브트리, 오른쪽 서브트리의 높이의 값을 확인하고 비교
  • 높이 차이를 초과하면, False
  • 초과하지 않는다면, 현재 왼쪽 서브트리, 오른쪽 서브트리의 높이에 + 1한 값을 현재 노드의 높이로 heights 배열에 저장

시간복잡도, 공간복잡도

• 시간 복잡도: 모든 노드에 한번씩 방문 O(N)
• 공간 복잡도: 스택에 저장되는 노드의 최대 개수 = 높이 O(H)

코드

class Solution:
    def isBalanced(self, root: Optional[TreeNode]) -> bool:
        if not root:
            return True

        stack = [(root, False)]
        heights = {}

        while stack:
            node, visited = stack.pop()

            if not node:
                continue

            if visited:
                leftHeight = heights.get(node.left, 0)
                rightHeight = heights.get(node.right, 0)

                if abs(leftHeight - rightHeight) > 1:
                    return False

                heights[node] = max(leftHeight, rightHeight) + 1
            else:
                stack.append((node, True))
                stack.append((node.left, False))
                stack.append((node.right, False))
        return True