Python으로 균형 이진 트리(Balanced Binary Tree) 판별하기

이진 트리(binary tree)에서 각 노드는 왼쪽 자식과 오른쪽 자식, 두 개의 자식 노드를 가질 수 있습니다. 어떤 이진 트리가 주어졌을 때 그 트리가 균형 상태인지 확인해야 하는 경우가 종종 있습니다. 왼쪽 서브트리와 오른쪽 서브트리의 높이 차이가 '1' 이하일 때, 그 이진 트리를 균형 트리(height-balanced tree)라고 합니다.

예시

입력-1

출력

True

설명

주어진 이진 트리는 [1, 2, 3, NULL, NULL, 6, 7]입니다. 왼쪽 서브트리와 오른쪽 서브트리의 높이 차이가 '1'이므로 높이 균형 트리입니다.

입력-2

출력

False

설명

주어진 이진 트리는 [1, 2, 3, 4, NULL, NULL, NULL, 5]입니다. 왼쪽 서브트리와 오른쪽 서브트리의 높이 차이가 '1'보다 크므로 높이 균형 트리가 아닙니다.

문제 해결 접근 방식

이 문제는 재귀적으로 해결할 수 있습니다. 먼저 왼쪽 서브트리와 오른쪽 서브트리의 높이를 각각 구한 뒤, 두 높이의 차이가 1 이하인지 확인합니다. 조건을 만족하면 True, 그렇지 않으면 False를 반환하고, 이 검사를 트리의 모든 노드에 대해 재귀적으로 수행합니다.

• 이진 트리의 노드들을 입력받습니다.
• 트리의 높이를 계산하는 함수를 정의합니다.
• 왼쪽 서브트리와 오른쪽 서브트리의 높이 차이가 '1'을 넘지 않는지 재귀적으로 검사하는 불리언 함수를 작성하고, 조건을 만족하면 True를 반환합니다.
• 최종 결과를 반환합니다.

구현 예제

class TreeNode:
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None

# 서브트리의 높이를 계산하는 함수
def get_height(node):
    if node is None:
        return 0
    return 1 + max(get_height(node.left), get_height(node.right))

# 트리가 균형인지 확인하는 함수
def is_balanced(root):
    if root is None:
        return True
    left_height = get_height(root.left)
    right_height = get_height(root.right)
    return (
        abs(left_height - right_height) <= 1
        and is_balanced(root.left)
        and is_balanced(root.right)
    )

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.right.left = TreeNode(6)
root.right.right = TreeNode(7)

if is_balanced(root):
    print("Balanced")
else:
    print("Not Balanced")

위 코드를 실행하면 아래와 같은 결과가 출력됩니다.

출력

Balanced

주어진 이진 트리 [1, 2, 3, NULL, NULL, 6, 7]은 왼쪽 서브트리와 오른쪽 서브트리의 높이 차이가 '1'이므로 높이 균형 트리입니다.

더 효율적인 방법: O(n) 최적화

위 방식은 노드마다 매번 높이를 새로 계산하므로 트리가 커지면 비효율적일 수 있습니다(최악의 경우 O(n²)). 다음과 같이 리프 노드 쪽부터 높이를 반환하면서 동시에 균형 여부를 함께 판단하면, 단 한 번의 순회(O(n))만으로 문제를 해결할 수 있습니다.

def is_balanced_optimized(root):
    def check(node):
        # 불균형이면 -1, 균형이면 해당 서브트리의 높이를 반환
        if node is None:
            return 0
        left = check(node.left)
        if left == -1:
            return -1
        right = check(node.right)
        if right == -1:
            return -1
        if abs(left - right) > 1:
            return -1
        return max(left, right) + 1

    return check(root) != -1

두 방법 모두 같은 결과를 반환하지만, 트리의 크기가 큰 실무 환경에서는 후자의 O(n) 방식을 사용하는 것이 좋습니다.