[리트코드] Balanced Binary Tree(python)

리트코드, Balanced Binary Tree_균형 이진 트리

TL;DR

• 트리(Tree)
• 깊이 우선 탐색(DFS)

문제 요약

1. 주어진 이진 트리가 높이 균형이 맞는지 아닌지를 판별하는 프로그램을 작성하라.
2. 높이 균형이 맞는 이진 트리는 모든 노드들의 왼쪽 및 오른쪽 서브 트리의 높이 차이가 1을 초과하지 않는 트리를 의미한다.

 

- 형제 노드 간의 차이가 1을 초과하지 않을 때는 True를, 그렇지 않은 경우에는 False를 반환하면 된다.

입출력 형태

입출력 예시

 

- 예시 1의 경우 모든 서브트리에서 높이 차이가 최대 1이다. 그렇기 때문에 True를 반환하였다.

- 예시 2의 경우 가장 하위에 있는 서브트리와([3, 4, 4])가 루트의 서브트리([1, 2, 2])와 높이 차이가 1을 초과하기 때문에 false를 반환하였다.

풀이

# 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 isBalanced(self, root: Optional[TreeNode]) -> bool:
        if not root:
            return True
        
        def dfs(node: TreeNode) -> int:
            if not node:
                return 0
            
            left = dfs(node.left)
            right = dfs(node.right)

            if left == -1 or right == -1 or abs(left - right) > 1:
                return -1
            
            return max(left, right) + 1
        
        return dfs(root) != -1

 

트리에서 DFS를 사용할 경우 각 노드에 대한 값을 할당하도록 작성할 수 있다.

서브트리를 만들 수 없는 리프 노드에 대해서는 0의 값을 반환하도록 할당한다. 그렇게 되면 리프 노드부터 1씩 값이 증가하는 형태로 값을 할당할 수 있다. 각 예시에 대해서 위의 코드를 실행하면 각 노드에는 그림과 같이 값이 할당되게 된다.

 

DFS에 의해 값이 할당되는 구조

 

자식 노드가 없는 노드들에 대해서는 0이 할당되고 리프노드부터 1이 할당되고 부모  노드는 자식 노드의 값 중 큰 값에 1을 더한 형태로 값이 할당되게 된다.

높이가 균형을 이루는 이진 트리의 경우 형제 노드의 값의 차가 1 이하인 경우라고 말할 수 있다.

그림에서 왼쪽인 예시 1번의 경우 형제 노드 간 값의 차가 최대 1이기 때문에 높이 균형 이진 트리라고 할 수 있다.

그림에서 오른쪽인 예시 2번의 경우 형제 노드 간 값의 차가 2인 서브 트리가 존재하기 때문에 높이 균형 이진 트리라고 할 수 없다.

 

이를 코드에서 구현한 부분이 dfs 함수 내 조건문이다.

if left == -1 or right == -1 or abs(left - right) > 1:
	return -1

 

이 조건문에 따라 코드를 실행하게 되면 형제 노드 간 값의 차가 1을 초과하는 순간 해당 노드에는 -1이 할당된다. -1이 할당된 노드가 존재한다는 것은 높이 균형 이진 트리를 구성할 수 없다는 뜻이다. 왼쪽 자식 노드와 오른쪽 자식 노드의 값이 하나라도 -1이라면 해당 이진 트리는 높이 균형 이진 트리가 될 수 없음을 의미한다. 따라서 조건문에 해당 조건인 `left == -1 or right == -1`이 추가되어야 한다.

 

이렇게 작성한 dfs 함수를 입력주어 주어진 이진 트리 root에서부터 실행하게 되면 그 결과는 정수 값을 반환한다. 이 때, 반환되는 값이 -1이라면 해당 이진 트리는 높이 균형을 이룰 수 없는 트리이고, 그렇지 않은 경우에는 높이 균형을 이루는 이진 트리이다.