[Swift] 백준 5639번: 이진 검색 트리

Lee Youjin·2023년 7월 1일

Swift로 백준 풀기

목록 보기

6/12

문제

이진 검색 트리는 다음과 같은 세 가지 조건을 만족하는 이진 트리이다.

노드의 왼쪽 서브트리에 있는 모든 노드의 키는 노드의 키보다 작다.

노드의 오른쪽 서브트리에 있는 모든 노드의 키는 노드의 키보다 크다.

왼쪽, 오른쪽 서브트리도 이진 검색 트리이다.

전위 순회 (루트-왼쪽-오른쪽)은 루트를 방문하고, 왼쪽 서브트리, 오른쪽 서브 트리를 순서대로 방문하면서 노드의 키를 출력한다. 후위 순회 (왼쪽-오른쪽-루트)는 왼쪽 서브트리, 오른쪽 서브트리, 루트 노드 순서대로 키를 출력한다.
예를 들어, 위의 이진 검색 트리의 전위 순회 결과는 50 30 24 5 28 45 98 52 60 이고, 후위 순회 결과는 5 28 24 45 30 60 52 98 50 이다.
이진 검색 트리를 전위 순회한 결과가 주어졌을 때, 이 트리를 후위 순회한 결과를 구하는 프로그램을 작성하시오.

입력

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

출력

입력으로 주어진 이진 검색 트리를 후위 순회한 결과를 한 줄에 하나씩 출력한다.

접근

처음에는 전위 순회한 결과를 바탕으로 트리를 만들고, 그대로 후위 순회해 출력하려고 했다. 근데 시간초과가 나버림 ... 아래는 시간초과가 난 코드다.

//
//  main.swift
//  ProblemSolving_5639
//
//  Created by YOUJIM on 2023/07/02.
//
import Foundation

class Node<T: Comparable> {
    let parent: T
    var leftChild: Node?
    var rightChild: Node?
    
    init(_ data: T) { self.parent = data }
}

class BinaryTree<T: Comparable> {
    var root: Node<T>?
    
    func add(_ data: T) {
        guard let root = self.root else { return self.root = Node.init(data) }
        var currentNode = root
        
        while true {
            if currentNode.parent > data {
                guard let nextNode = currentNode.leftChild else { return currentNode.leftChild = Node.init(data) }
                currentNode = nextNode
            }
            else {
                guard let nextNode = currentNode.rightChild else { return currentNode.rightChild = Node.init(data) }
                currentNode = nextNode
            }
        }
    }
    
    func postOrder(_ node: Node<T>?) {
        guard let node = node else { return }
        
        BinaryTree().postOrder(node.leftChild)
        BinaryTree().postOrder(node.rightChild)
        print(node.parent)
    }
}

var binaryTree = BinaryTree<Int>()
while let line = readLine() { binaryTree.add(Int(line)!) }
binaryTree.postOrder(binaryTree.root)

아니 도대체 왜? 시간 초과가? 났지? 라는 생각이 들긴 했다. 근데 생각해보면 이렇게 풀리면 골드 문제가 아닐 것 같다는 생각도 들고 ... 뭔가 전위 순회한 결과를 입력으로 준 이유가 있을 것 같았다. 귀찮아서 그냥 트리를 만들긴 했지만 ㅜㅜ

전위 순회는 루트 노드 -> 루트 노드의 왼쪽 노드 -> 루트 노드의 오른쪽 노드 방향으로 순회하는 방법이다.
그렇다는 말은 입력의 가장 첫 번째는 루트 노드의 값이 되고, 이후 나오는 값은 루트 노드의 바로 왼쪽 노드, 그 다음은 왼쪽 노드의 바로 왼쪽 노드 ... 가 된다. 이 말은 나오는 값이 계속 작아진다는 말이다. 예시로 문제의 입력 값을 보면 다음과 같다.

50
30
24
5
28
45
98
52
60

트리의 루트 노드인 50에서 28이 나올 때까지 입력값이 계속 감소한다. 이 값들은 트리의 가장 왼쪽에 있는 값이 된다.

그렇다면 값을 계속 입력받다가 바로 전에 입력받은 값보다 이번에 입력받은 값이 크다면? 왼쪽 노드에서 오른쪽 노드로 옮겨갔다는 것을 뜻한다. 이러한 과정을 재귀적으로 계속 반복해주면 된다.

... 라고 자신있게 말했지만 구현은 힘들었다. 왜 머리에는 있는데 구현이 안되지? 결국 질문 게시판을 들어가봤는데 나랑 똑같은 의문을 가지신 분이 계셨다!
같은 시간복잡도가 쓰이는데 하나는 통과, 하나는 시간초과인 이유?
이 분도 처음에 나처럼 트리를 구현하시고 시간 초과가 나신 케이스였다. 답변을 끌고 와보면

시간 복잡도는 점근적인 효율성을 보여줄 뿐입니다. 단순히 통과/시간 초과로만 나누지 말고, 통과된 코드가 얼마나 시간이 걸렸는지 확인해보세요. 이 문제의 파이썬의 제한은 5초인데, 통과된 코드도 3.6초나 걸렸습니다. 즉, 같은 시간 복잡도라도 세부 단계의 효율성이 1.5배만 안 좋아도 시간이 초과될 수 있다는 뜻입니다. 객체라는 건 그 자체로 상당히 무거운 녀석입니다. 이것들을 전부 노드 하나당 하나의 객체를 할당하고 서로 엮어서 실제로 트리로 구현하는 건 대충 보기에도 매우 무거워 보입니다.

라고 한다. 조만간 코드별로 걸리는 시간 알아보는 법을 좀 공부해야겠다.
그리고 결국 참고해서 풀었는데 생각보다 접근 파트가 좀 복잡해져서 ㅜㅜ 이번에는 코드 파트에서 설명하겠다!

그리고 이 문제에서 주의할 점은 딱히 입력 개수를 주거나 입력이 끝났음을 알려주는 게 아니기 때문에 while 문을 이용해 입력이 끝나면 EOF 로 마무리해줘야 한다는 점이다. 난 이런 문제를 처음 풀어봐서 조금 해멨는데, 이런 문제에서는 다음과 같이 입력을 받으면 된다.

while let line = readLine() { // 실행할 구문 }

이 코드에서 line 의 값이 nil 이 아니면 중괄호 안 구문이 실행된다. control + D 를 해주면 EOF가 실행된다. 만약 while true 로 조건을 설정하고 중괄호 안에서 입력받을 경우 백준에서 런타임 에러가 날 수 있다고 하니 위처럼 작성하는 게 좋을 것 같다.
[Swift] Swift로 입력 받기 (EOF, 무한 입력 받기 해결 성공일기!)
[Swift] While EOF / 백준에서 while 입력 끝 날때까지 받기 / Swift 끝까지 입력 받기

코드

import Foundation

var binaryTree = [Int]()

func printBinaryTree(_ first: Int, _ end: Int) {
    if first > end { return }
    var mid = end + 1 // 모든 원소가 루트 노드보다 작을 경우 29번째 줄에서 printBinaryTree(0, end)가 됨
                      // 그러면 25번째 줄에서 오른쪽 노드를 조사할 때 printBinaryTree(end + 1, end)가 되므로 return되고, 이는 오른쪽 노드가 없다는 것을 뜻함
    
    for i in (first + 1)..<(end + 1) {
        if binaryTree[first] < binaryTree[i] { // 루트 노드보다 큰 원소가 나올 경우
            mid = i // 왼쪽 서브트리와 오른쪽 서브트리를 나누는 지점 할당
            break
        }
    }
    
    printBinaryTree(first + 1, mid - 1) // 왼쪽 노드 조사
    printBinaryTree(mid, end) // 오른쪽 노드 조사
    print(binaryTree[first])
}

while let line = readLine() { binaryTree.append(Int(line)!) }
printBinaryTree(0, binaryTree.count - 1)

접근에서 한참 헤맨 것 치고는 코드가 간단하다.

일단 printBinaryTree() 는 전위 순회로 입력받은 트리를 조사하고 후위 순회대로 출력해주는 함수다.
파라미터로 입력받은 first 와 end 는 배열의 인덱스를 뜻하는데, 이후 mid 를 통해 왼쪽 서브트리와 오른쪽 서브트리를 나누는 기준을 설정한다. for 문을 이용해 first 인덱스부터 end 인덱스까지 돌면서 만약 루트 노드(first)보다 더 큰 값이 나온다면 그 부분(코드에서는 i)이 왼쪽 서브트리와 오른쪽 서브트리를 나누는 기준이 되므로, mid에 i를 할당해준다. 이후 재귀를 이용해 왼쪽 서브트리와 오른쪽 서브트리를 각각 조사해주고 루트 노트를 출력한다.

그리고 주석으로도 적어놨지만 왜 처음에 mid 에 end + 1 을 할당하냐면 모든 원소가 루트 노드보다 작을 경우를 고려해야 하기 때문이다. 자세한 설명은 코드에 주석으로 적어놔서 생략하고 넘어가겠다!