[백준(파이썬/Python)] 5639_이진 검색 트리

kindof

·

2021. 8. 8. 16:00

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

 

5639번: 이진 검색 트리

트리를 전위 순회한 결과가 주어진다. 노드에 들어있는 키의 값은 106보다 작은 양의 정수이다. 모든 값은 한 줄에 하나씩 주어지며, 노드의 수는 10,000개 이하이다. 같은 키를 가지는 노드는 없다

www.acmicpc.net

 

 

이 문제는 시간초과 때문에 고통받았던 문제입니다.

 

아래는 기존의 풀이이고 이렇게 하면 (전위순회를 통해 트리를 생성하는 시간) + (후위순회를 하는 시간(덤으로 재귀까지))때문에 시간초과가 발생합니다.

 

기본적인 트리 순회는 스택이나 재귀를 이용해서 하면 되지만 이 문제에서는 효율적인 탐색을 요구한 것이죠.

 

// 시간초과

import sys
sys.setrecursionlimit(10 ** 6)


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

class BinaryTree():
    def __init__(self):
        self.root = None;

    def insert(self, newNode):
        if self.root == None:
            self.root = newNode        

        else:
            self.currNode = self.root
            while True:
                if newNode.data < self.currNode.data:
                    if self.currNode.left != None:
                        self.currNode = self.currNode.left
                    else:
                        self.currNode.left = newNode
                        break
                else:
                    if self.currNode.right != None:
                        self.currNode = self.currNode.right
                    else:
                        self.currNode.right = newNode
                        break

    def postOrder(self, root):
        result, stack = [], []
        stack.append(root)

        while stack:
            curr = stack.pop()
            result.append(curr)

            if curr.left:
                stack.append(curr.left)
            if curr.right:
                stack.append(curr.right)
        
        while(result):
            print(result.pop().data)



tree = BinaryTree();

while True:
    try: 
        data = int(sys.stdin.readline())
        node = Node(data)
        tree.insert(node)
    except:
        break
tree.postOrder(tree.root)

 

그래서 아래와 같이 이진 탐색을 이용해서 문제를 해결해야 합니다.

 

전위 순회를 한 결과는 [root, {root보다 작은 값들}, {root보다 큰 값들}] 꼴이기 때문에 저기 {root보다 작은 값들} -> {root보다 큰 값들} -> root 순서로 탐색을 진행해주면 됩니다.

 

이 때 각 구간의 인덱스를 찾기 위해서 시작 지점을 root로 지정해주고 그것보다 큰 값이 나타나는 구간을 기준으로 짤라주면 됩니다.

 

 

// 통과 코드(참고)

import sys
sys.setrecursionlimit(10 ** 6)

def postOrder(start, end):
    if start > end:
        return
    
    root = preOrder[start]
    idx = start + 1

    while idx <= end:
        if preOrder[idx] > root:
            break
        idx += 1
    
    postOrder(start +1 , idx-1)
    postOrder(idx, end)
    print(root)

preOrder = []
while True:
    try:
        preOrder.append(int(sys.stdin.readline()))
    except:
        break
postOrder(0, len(preOrder) - 1)