[백준 1991]python 인생 처음 이진 트리 구현

오늘은 백준 1991 트리순회(이진트리)를 풀었다.
문제 링크는 아래에 있다.
https://www.acmicpc.net/problem/1991

1991번: 트리 순회

 

백준 1991 예제 이진트리 사진

 
기본적으로 문제가 원하는 것은 이진트리를 넣어줄 테니  전위 순회(preorder traversal), 중위 순회(inorder traversal), 후위 순회(postorder traversal)를 해달라는 문제였다.
맨 처음 문제를 보고 드는 생각은 "뭔 소리지?" 였다.
어제까지만 해도 dfs문제를 풀고 원숭이처럼 기뻐하는 내게 너무 큰걸 기대하는게 아닌가...
혼자 끙끙거리다 같은 동기형에게 물어보니 다행이 친절히 설명해 주었다.
 

---

 
우선 전위 순회(preorder traversal)

문제에 있던 (루트) -> (왼쪽 자식) -> (오른쪽 자식) 이 힌트였다.
우선 가장 꼭대기인 "A"에서 시작하자. (사진을 보면서 따라오면 이해가 잘된다)
위에 있는 순서대로 (루트)부터 시작한다.
즉 "A"가 들어있는 노드(동그라미라고 생각해도 된다)를 탐색한다. 따라서 맨 처음 탐색은 A가 출력된다.
다름으로는 (왼쪽 자식) 으로 넘어간다.
"B"가 들어있는 노드로 들어와서 해야 할 것은 (루트) 이다. ( (루트) -> (왼쪽 자식) -> (오른쪽 자식) 다시 반복)
따라서 "B"가 탐색으로 출력된다.
다음으로 (왼쪽 자식)으로 넘어간다. ( "B"노드에서의 다음 명령은 (오른쪽 자식) 이다)
"D"가 들어있는 노드로 들어온다,.
우선 (루트)를 해서 "D"를 출력한다.
그 이후 (왼쪽 자식)을 가야하는데 왼쪽자식이 없다.
그렇다면 이 명령의 다음인 (오른쪽 자식)으로 넘어간다.
이때 오른쪽 자식도 없다. 그렇다면 "D"가 들어있는 노드에서 할 일은 끝난 것이다.
그럼 그 다음 명령은? "B"가 들어있는 노드에서의 (오른쪽 자식) 명령이다.
이때에도 이것도 없기 때문에 그다음 명령인 "A"가 들어있는 노드에서 (오른쪽 자식) 명령을 실행한다.
"C"노드로 들어왔다. 다시 (루트) 명령부터 시작해 "C"를 출력한다.
다음으로 (왼쪽 자식)을 실행해 "E"가 들어있는 노드로 들어간다.
"E"가 들어있는 노드에서 다시 (루트)를 실행해 "E"를 출력한다.
다음 명령인 (왼쪽 자식)도 불가능하고 (오른쪽 자식)도 불가능해서 다음 명령인 "C"노드에서 (오른쪽 자식)을 실행한다.
이후 "F"와 "G"는 위와 같이 생각하면 순서대로 출력된다.
 
위의 말을 잘 따라왔다면 중위 순회(inorder traversal)와 후위 순회(postorder traversal)는 각각 
(왼쪽 자식) (루트) (오른쪽 자식), (왼쪽 자식) (오른쪽 자식) (루트) 이라는 것만 알면 할 수 있다.
 

---

 
1. 리스트로 이진 트리 구현
맨 처음에는 리스트로 이진 트리를 구현해 보았다.
에초에 이진 트리에 대해서는 힙을 구현할 때 보았던 리스트밖에 없었다.
기본 원리는 이렇다.
이진 트리를 구현한 arr를 하나 만들어준다.
arr에서 인덱스가 1(a)인 방에 이진 트리의 꼭대기 노드의 값을 넣어준다.
이때 노드의 왼쪽과 오른쪽에 있는 값은 각각 1*2, 1*2 + 1인 인덱스 방에 넣어준다.(a * 2, a * 2 + 1)
모든 노드에 대해 (자신의 인덱스 번호 * 2)에 왼쪽노드값을, (자신의 인덱스 번호 * 2 + 1)에 오른쪽 노드 값을 넣으면 된다.
예제를 구현해보면

 idx = 0   1    2    3    4   5   6    7   8  9  10 11 12 13 14  15
arr = [0, 'A', 'B', 'C', 'D', 0, 'E', 'F', 0, 0, 0, 0, 0, 0, 0, 'G']

이렇게 된다. (없는 노드는 0을 넣어준다)
파이썬으로 문제를 풀어보면

n = int(input())
arr = [0, 0]
arr[1] = "A"

for _ in range(n):
    a, b, c = input().split(" ")
    sub2 = arr.index(a) * 2 + 1
    sub = len(arr) - 1
    if sub < sub2:
        arr += [0] * (sub2 - sub)
    if b != ".":
        arr[sub2 - 1] = b
    if c != ".":
        arr[sub2] = c

def preorder(now): # 자기 자신, 왼쪽, 오른쪽
    print(arr[now], end = "")
    if arr[now * 2] != 0: preorder(now * 2)
    if arr[now * 2 + 1] != 0: preorder(now * 2 + 1)

def inorder(now): # 왼쪽, 자기자신, 오른쪽
    if arr[now * 2] != 0: inorder(now * 2)
    print(arr[now], end = "")
    if arr[now * 2 + 1] != 0: inorder(now * 2 + 1)

def postorder(now): # 왼쪽 오른쪽, 자기자신
    if arr[now * 2] != 0: postorder(now * 2)
    if arr[now * 2 + 1] != 0: postorder(now * 2 + 1)
    print(arr[now], end = "")

preorder(1)
print("")
inorder(1)
print("")
postorder(1)

이렇게 된다.
처음으로 이진 트리를 만들어 봐서 싱글벙글 하고 있었는데 백준이 나에게 준 축하선물은 "메모리 초과"였다. (으어어ㅓ)
근데 메모리 초과가 뜰거라고 어느정도 생각은 하고있었다.
에초에 왼쪽 아래로만 자식노드가 계속 생기면 배열의 길이는 2**25인데 메모리 초과가 안뜨는게 이상한거다.
솔직히 여기서부터는 모르겠어서 구글에 파이썬으로 이진트리 구현 을 검색했다.
그리고 다른 사람이 class로 만든 코드를 보면서 나만의 방식으로 만들어봤다.
 
2. class로 이진 트리 구현
class로 노드 구현은 할만 했다.
사실 이진 트리 class도 만들만은 했다.
문제는 이진 트리에서 노드를 추가할 때 어따가 추가할지 모른다는 점이었다.
이 점에 대해 나 혼자 머리 싸매고 있을 때, 번득이는 아이디어가 떠올랐다.
"dfs로 제귀를 써서 트리를 탐색해보면 되지 않을까?"
바로 class 안에 제귀를 사용해서 탐색 해 보았다.
결과는 대성공. (씨익)
아래는 class로 만든 이진 트리로 만든 나만의 백준 답지다.

# class로 구현한 이진트리
n = int(input())
class Node:
    def __init__(self, item):
        self.item = item
        self.left = None
        self.right = None
        
class TwoLoseTree:
    def __init__(self, item):
        self.first = item # Node가 item임
    def add(self, now, up, datat, lr): # lr 0은 왼쪽, 1은 오른쪽, up은 위 노드 데이터, datat는 데이터
        if now.item == up:
            if lr == 0:
                now.left = Node(datat)
            else:
                now.right = Node(datat)
        if now.left != None:
            self.add(now.left, up, datat, lr)
        if now.right != None:
            self.add(now.right, up, datat, lr)
        
tree = TwoLoseTree(Node("A"))
for _ in range(n):
    a, b, c = input().split(" ")
    if b != ".":
        tree.add(tree.first, a, b, 0)
    if c != ".":
        tree.add(tree.first, a, c, 1)

def preorder(now): # 자기 자신, 왼쪽, 오른쪽
    print(now.item, end = "")
    if now.left != None: preorder(now.left)
    if now.right != None: preorder(now.right)

def inorder(now): # 왼쪽, 자기자신, 오른쪽
    if now.left != None: preorder(now.left)
    print(now.item, end = "")
    if now.right != None: preorder(now.right)

def postorder(now): # 왼쪽 오른쪽, 자기자신
    if now.left != None: preorder(now.left)
    if now.right != None: preorder(now.right)
    print(now.item, end = "")

preorder(tree.first)
print("")
inorder(tree.first)
print("")
postorder(tree.first)

 
코드를 만들고 여러가지 예제를 돌려보고 스스로 뿌듯했다.
실제로 돌려보니 44ms로 만족스러운 속도를 보여줬다.
비록 class로 구현하다 보니 트리 내부를 보기 힘든점이 있긴 했지만 메모리나 시간이 상당히 줄어든 것을 보고 만족했다.
이렇게 넘어갈 찰나, 다른 사람들의 통과한 코드를 보다가 놀라온 코드르 발견했다.
 
3. 딕셔너리로 이진 트리 구현
세상에 간단하면서 속도도 빠른 코드를 발견해 버렸다.
바로 딕셔너리 키를 특정 노드로, value 값을 노드의 자식들로 만들어 버린 것이다.
이런 좋은 코드는 놓칠수 없지 나 혼자 다시 딕셔너리로 구현해 보았다.
아래는 딕셔너리로 구현한 내 코드다.

# 딕셔너리로 만들기.
dic = {}
n = int(input())
for _ in range(n):
    a, b, c = input().split(" ")
    dic[a] = (b, c)

def preorder(now):
    print(now, end = "")
    if dic[now][0] != ".": preorder(dic[now][0])
    if dic[now][1] != ".": preorder(dic[now][1])

def inorder(now):
    if dic[now][0] != ".": inorder(dic[now][0])
    print(now, end = "")
    if dic[now][1] != ".": inorder(dic[now][1])
        
def postorder(now):
    if dic[now][0] != ".": postorder(dic[now][0])
    if dic[now][1] != ".": postorder(dic[now][1])
    print(now, end = "")

preorder("A")
print("")
inorder("A")
print("")
postorder("A")

딕셔너리로 만들다 보니 한번에 트리 내부를 보기도 쉽고, 단번에 키값으로 트리에 들어가기 때문에 빠르고, 무려 이진트리가 아닌 다른 트리도 이렇게 구현 가능하기 때문에 매우 좋은 코드라 생각한다.
 

---

 
이렇게 이진 트리를 배열, 클래스, 딕셔너리로 구현 해 보았다.
지금까지는 배열 방법밖에 몰랐는데 클래스나 딕셔너리 방법을 익혀서 매우 유익한 문제였다.
다음에는 전위 순회, 중위 순회, 후위 순회가 무엇인지, 어디에 쓰이는 지에 대해 알아봐야겠다.
 

---

 
reference

• http://ejklike.github.io/2018/01/09/traversing-a-binary-tree-1.html
• https://wikidocs.net/16