이진 트리 균형 확인

이진 트리가 균형 잡혀있는지 확인하는 함수를 작성하라. 균형 잡힌 트리란 모든 노드에 대해서 왼쪽 부분 트리의 높이와 오른쪽 부분 트리의 높이의 차이가 최대 하나인 트리를 의미한다.

이 문제를 보자마자 드는 생각을 정리해보자면 대충 이렇다.

1. 왼쪽 / 오른쪽 서브 트리의 높이 구하기.

2. 높이 차 구하기.

2-1. 1이상일 경우 False 반환.

2-2. 그렇지 않을 경우 subtree의 높이 탐색.

이를 코드로 구현해보면 다음과 같다.

def getHeight(root):
 
    if not root :
        return -1
 
    return max(getHeight(root.left), getHeight(root.right))
 
def isBalanced(root):
 
    if not root :
        return True
 
    diff = getHeight(root.left)-getHeight(root.right)
 
    if abs(diff) > 1 :
        return False
    else :
        return isBalanced(root.left) and isBalanced(root.right)
 

하지만, 별로 효율적인 방법은 아닌것 같다. 각 노드에서 전체 하위 트리를 재귀적으로 탐색하므로, 같은 노드에 대해 getHeight가 반복적으로 호출된다. 이 알고리즘은 \(O(nlogn)\)인데, 상위 노드들이 하위 노드를 전부 한 번씩 건드리기 때문이다. 따라서 getHeight 호출 횟수를 줄일 필요가 있다. 사실 높이를 구하는 과정에서도 균형이 잡혀있는지 검사를 할 수 있다. 하위 트리가 균형 잡혀 있지 않다면 그냥 에러를 반환하면 된다. 그럼 getHeight 함수를 변경해보자.

def getHeight(root):
    
    if not root :
        return
    
    lHeight = getHeight(root.left)
    if lHeight == float("inf") :
        return float("inf")
 
    rHeight = getHeight(root.right)
    if rHeight == float("inf") :
        return float("inf")
 
    diff = lHeight - rHeight :
    if abs(diff) > 1:
        return float("inf")
    else :
        return max(lHeight, rHeight) + 1
 
def isBalanced(root):
 
    return getHeight(root) != float("inf")

 

이 코드의 시간 복잡도는 \(O(N)\)이고 공간 복잡도는 \(O(H)\)이다. 여기서 H는 트리의 높이를 나타낸다.

출처 : 게일 라크만 맥도웰. Cracking the coding interview. 프로그래밍 인사이트.