BOJ-11725 트리의 부모 찾기 Swift

문제

루트 없는 트리가 주어진다. 이때, 트리의 루트를 1이라고 정했을 때, 각 노드의 부모를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 노드의 개수 N (2 ≤ N ≤ 100,000)이 주어진다. 둘째 줄부터 N-1개의 줄에 트리 상에서 연결된 두 정점이 주어진다.

출력

첫째 줄부터 N-1개의 줄에 각 노드의 부모 노드 번호를 2번 노드부터 순서대로 출력한다.

내가 푼 풀이

- 주어진 연결된 두 정점을 연결리스트로 저장한다.

- 트리의 루트는 1 이므로 1부터 인접한 곳을 탐색한다.

- 이때 인접한 곳의 부모노드는 이전에 탐색된 노드의 번호가 되는 것이다.

- 인접한 노드 중 이미 탐색한 곳을 제외하고 새롭게 인접한 곳을 탐색한다.

- 부모노드를 배열에 따로 저장하고 출력

 

코드로 구현하면 다음과 같다.

import Foundation

// 부모노드배열, 트리 연결리스트
let count = Int(readLine()!)!
var tree = [Int: [Int]]()
var parent = Array(repeating: 0, count: count+1)

// 연결리스트 입력
for i in 0..<count-1 {
    let input = readLine()!.split(separator: " ").map{ Int($0)! }
    if tree[input[0]] == nil {
        tree[input[0]] = [input[1]]
    } else {
        tree[input[0]]!.append(input[1])
    }
    if tree[input[1]] == nil {
        tree[input[1]] = [input[0]]
    } else {
        tree[input[1]]!.append(input[0])
    }
}

// 시작점은 루트노드 1
var needVisitQueue = [1]

// BFS
// 인접한 노드 탐색
// 이미 탐색한 곳은 제외하고 새롭게 탐색한다.
// 인접한 노드의 부모노드는 이전에 탐색된 노드번호가 된다.
while !needVisitQueue.isEmpty {
    var node = needVisitQueue.removeLast()
    var arr = [Int]()
    for i in tree[node]! {
        if parent[i] == 0 && i != node {
            parent[i] = node
            arr.append(i)
        }
    }
    needVisitQueue += arr
}

// 부모노드 출력
for i in 2...count {
    print(parent[i])
}