백준 트리 순회 코드 및 해설 (파이썬)

2021. 6. 8. 12:15algorithm

반응형

https://www.acmicpc.net/problem/1991

 

1991번: 트리 순회

첫째 줄에는 이진 트리의 노드의 개수 N(1≤N≤26)이 주어진다. 둘째 줄부터 N개의 줄에 걸쳐 각 노드와 그의 왼쪽 자식 노드, 오른쪽 자식 노드가 주어진다. 노드의 이름은 A부터 차례대로 영문자

www.acmicpc.net

 

트리 구조를 TreeNode라는 클래스로, 리스트를 입력으로 받아 트리 구조로 만들어주는 함수 makeTree를 정의했습니다.
우선 입력을 받아와 노드의 개수 N을 제외하고 3개씩 묶어 nodes라는 리스트에 저장합니다. 이후 makeTree 함수를 이용해 트리 구조로 바꿔줍니다.
트리 구조로 바꿔준 input을 순서대로 전위순회 preorder, 중위순회 inorder, 후위순회 postorder 함수에 전달합니다.
전위순회는 먼저 루트를 result에 추가하고 왼쪽 자식, 오른쪽 자식을 차례대로 다시 preorder 함수에 recursive하게 전달합니다.
중위순회는 왼쪽 자식을 inorder 함수에 recursive하게 전달한 후, 루트를 result에 추가하고, 마지막으로 오른쪽 자식을 inorder 함수에 recursive하게 전달합니다.
후위순회는 왼쪽 자식, 오른쪽 자식을 차례대로 postorder 함수에 recursive하게 전달합니다. 이후 루트를 result에 추가합니다.

 

import sys

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right
            
    
def makeTree(cur_node, arr):
    if cur_node[0] != '.':
        temp = TreeNode(cur_node[0])  
        root = temp
        
        cur_left_node = None
        cur_right_node = None
        
        for a in arr:
            if a[0] == cur_node[1]:
                cur_left_node = a 
            if a[0] == cur_node[2]:
                cur_right_node = a 
                
        if cur_left_node:
            root.left = makeTree(cur_left_node, arr)
        if cur_right_node:
            root.right = makeTree(cur_right_node, arr)
    else:
        root = None
        
    return root

# 입력 받아오기 
inputs = list(sys.stdin.read().split())
N = int(inputs[0])

nodes = []
for i in range(0,len(inputs[1:]),3):
    nodes.append(inputs[1:][i:i+3])

# 트리 구조로 만들기 
tree = makeTree(nodes[0], nodes)

# 전위순회
def preorder(tree, result):
    if tree:
        # 루트 
        result += tree.val
        # 왼쪽 자식, 오른쪽 자식
        children = [tree.left, tree.right]
        for child in children:
            if child:
                result = preorder(child, result)
    return result

# 중위순회
def inorder(tree, result):
    if tree:
        # 왼쪽 자식
        if tree.left:
            result = inorder(tree.left, result)
        # 루트
        result += tree.val
        # 오른쪽 자식 
        if tree.right:
            result = inorder(tree.right, result)
    return result

# 후위순회
def postorder(tree, result):
    if tree:
        # 왼쪽 자식, 오른쪽 자식 
        children = [tree.left, tree.right]
        for child in children:
            if child:
                result = postorder(child, result)
        # 루트 
        result += tree.val
    return result

print(preorder(tree, result=''))
print(inorder(tree, result=''))
print(postorder(tree, result=''))
반응형