파이썬으로 전위 순회와 후위 순회 결과에서 이진 트리 구성하기

문제 소개

전위 순회(preorder)와 후위 순회(postorder) 결과 두 개의 순회 순서가 주어졌을 때, 이 정보만으로 원래의 이진 트리를 복원하는 것이 목표입니다. 예를 들어 전위 순회가 [1,2,4,5,3,6,7], 후위 순회가 [4,5,2,6,7,3,1]이라면 아래와 같은 이진 트리가 생성됩니다.

참고로 전위 순회와 후위 순회만으로는 자식이 하나뿐인 노드가 왼쪽 자식인지 오른쪽 자식인지 판별할 수 없으므로, 이 방법으로 복원되는 트리는 항상 유일하지 않을 수 있다는 점을 알아두면 좋습니다.

해결 접근 방법

스택(stack)을 활용하면 두 순회 결과를 한 번의 순회로 비교하면서 효율적으로 트리를 구성할 수 있습니다. 단계별 과정은 다음과 같습니다.

pre[0] 값을 가진 루트 노드를 생성하고, 이 노드를 스택에 삽입합니다.
포인터 i는 1, j는 0으로 초기화합니다.
i가 pre의 길이보다 작고 j가 post의 길이보다 작은 동안 다음을 반복합니다.
- 스택 맨 위 노드의 값이 post[j]와 같다면 j를 1 증가시키고, 스택에서 pop한 뒤 다음 반복으로 넘어갑니다. 이는 해당 서브트리의 처리가 끝났음을 의미합니다.
- pre[i] 값을 가진 새 노드를 생성합니다.
- 스택 맨 위 노드의 왼쪽 자식이 비어 있으면 새 노드를 왼쪽 자식으로 연결하고, 그렇지 않으면 오른쪽 자식으로 연결합니다.
- 새 노드를 스택에 삽입하고 i를 1 증가시킵니다.
반복이 끝나면 루트 노드(ans)를 반환합니다.

구현 예제

아래는 파이썬으로 작성한 전체 구현 코드입니다. 트리 구성 로직과 함께 결과를 확인하기 위한 레벨 순회(level order) 출력 함수도 포함되어 있습니다.

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

def height(root):
    if root is None:
        return 0
    else:
        # 왼쪽과 오른쪽 서브트리의 높이를 계산
        l_height = height(root.left)
        r_height = height(root.right)
        # 더 큰 값에 1을 더해 반환
        if l_height > r_height:
            return l_height + 1
        else:
            return r_height + 1

def print_given_level(root, level):
    if root is None:
        return
    if level == 1:
        print(root.data, end=',')
    elif level > 1:
        print_given_level(root.left, level - 1)
        print_given_level(root.right, level - 1)

def level_order(root):
    print('[', end='')
    h = height(root)
    for i in range(1, h + 1):
        print_given_level(root, i)
    print(']')

class Solution(object):
    def constructFromPrePost(self, pre, post):
        ans = TreeNode(pre[0])
        stack = [ans]
        i = 1
        j = 0
        while i < len(pre) and j < len(post):
            if stack[-1].data == post[j]:
                j += 1
                stack.pop(-1)
                continue
            node = TreeNode(pre[i])
            if not stack[-1].left:
                stack[-1].left = node
            else:
                stack[-1].right = node
            stack.append(node)
            i += 1
        return ans

ob = Solution()
pre = [1,2,4,5,3,6,7]
post = [4,5,2,6,7,3,1]
tree = ob.constructFromPrePost(pre, post)
level_order(tree)

입력

pre = [1,2,4,5,3,6,7]
post = [4,5,2,6,7,3,1]

출력

[1,2,3,4,5,6,7]

복잡도 분석

이 알고리즘은 각 노드를 정확히 한 번씩 방문하므로 시간 복잡도는 O(n)입니다. 또한 최악의 경우(편향된 트리) 스택에 모든 노드가 저장될 수 있으므로 공간 복잡도 역시 O(n)입니다.