[Leetcode] 리트코드 - 전위, 중위 순회 결과로 이진 트리 구축(construct-binary-tree-from-preorder-and-inorder-traversal) 파이썬(python) 풀이

문제

트리의 전위, 중위 순회 결과를 입력값으로 받아 이진 트리를 구축하라.

전위 순회 결과: [3,9,20,15,7]
중위 순회 결과: [9,3,15,20,7]

결과는 밑의 그림과 같은 이진 트리가 된다.

코드

from typing import List


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


class Solution:
    def buildTree(self, preorder: List[int], inorder: List[int]) -> TreeNode:
        if inorder:
            # 전위 순회 결과는 중위 순회 분할 인덱스
            index = inorder.index(preorder.pop(0))

            # 중위 순회 결과 분할 정복
            node = TreeNode(inorder[index])
            node.left = self.buildTree(preorder, inorder[0:index])
            node.right = self.buildTree(preorder, inorder[index + 1:])


            return node

풀이

순회에는 크게 전위, 중위, 후위 순회가 있으며 이 셋 중 2가지만 있어도 이진 트리를 복원할 수 있다. 예제의 입력값이 너무 단순하므로, 트리 그림을 1부터 9까지 좀 더 복잡한 형태로 새로 구성해보자.

전위의 첫 번째 값은 부모 노드다. 즉 전위 순회의 첫 번째 결과는 정확히 중위 순회 결과를 왼쪽과 오른쪽으로 분할 시키는 역할을 한다. 중위 순회의 분할 정복 문제로 바꾸는 것이다. 두 번째로 왼쪽 노드의 2는 중위 순회 결과를 정확히 반으로 가르고, 각각 왼쪽 자식은 4, 오른쪽 자식은 5로 마무리 한다.

오른쪽의 경우 3이 첫 번째 값인데, 마침 중위 순회에서는 맨 오른쪽에 위치해 있다. 이 말은 3의 오른쪽 자식 노드는 존재하지 않는다는 얘기다. 이후 남아 있는 노드들을 이후에도 계속 분할을 시도하면, 이 그림처럼 트리 형태로 최종적으로 구성할 수 있다.

index = inorder.index(preorder.pop(0))

먼저, 전위 순회 첫 번째 결과를 가져와 중위 순회를 분할하는 인덱스로 한다.

node = TreeNode(inorder[index])
node.left = self.buildTree(preorder, inorder[0:index])
node.right = self.buildTree(preorder, inorder[index + 1:])

이 값을 현재 노드로 구성하고, 이를 기준으로 중위 순회 결과를 쪼개서 왼쪽, 오른쪽으로 각각 마무리 될 때 분할 정복 구조로 재귀 호출하면, 트리를 구성할 수 있다.

밑은 재귀 함수가 호출될 때마다 변수값 그림.

출처
파이썬 알고리즘 인터뷰 (글 : 박상길 그림 : 정진호) [책만]