Today I Learned

오늘은 백준 코테 공부!

트리의 부모 찾기

백준 11725번 https://www.acmicpc.net/problem/11725

문제

루트 없는 트리가 주어진다. 이때, 트리의 루트를 1이라고 정했을 때, 각 노드의 부모를 구하는 프로그램을 작성하시오. 첫째 줄에 노드의 개수 N (2 ≤ N ≤ 100,000)이 주어진다. 둘째 줄부터 N-1개의 줄에 트리 상에서 연결된 두 정점이 주어진다. 첫째 줄부터 N-1개의 줄에 각 노드의 부모 노드 번호를 2번 노드부터 순서대로 출력한다.

• 입력
  12
  1 2
  1 3
  2 4
  3 5
  3 6
  4 7
  4 8
  5 9
  5 10
  6 11
  6 12
• 출력
  1
  1
  2
  3
  3
  4
  4
  5
  5
  6
  6

문제 풀이

• DFS나 BFS같은 백트래킹을 이용해 트리를 탐색하면서 각 노드의 부모를 찾는 것이 맞는 방법이라고 생각해 DFS로 풀었다.

1. 각 노드의 부모를 저장하는 parent 배열과 각 노드의 연결된 노드를 저장하는 graph 2차원 배열을 초기화한다. repeat(n-1) 구문 안에서 그래프 정보를 입력받아 graph에 저장한다.

2. DFS 함수를 만든다. 먼저 현재 노드의 부모노드를 parent[node]에 저장한다. 그리고 해당 노드에 연결된 모든 노드들에대해 반복문을 수행한다. 만약 다음 노드가 부모 노드가 아니면 재귀적으로 DFS를 호출한다.

3. DFS의 결과가 parent에 저장되고 2~n까지 결과를 차례대로 출력하면 완료!

Kotlin 코드

fun main() {
    val n = readln().toInt()

    val parent = IntArray(n + 1)
    val graph = Array(n + 1) { ArrayList<Int>() }

    repeat(n - 1) {
        val (a, b) = readln().split(" ").map { it.toInt() }
        graph[a].add(b)
        graph[b].add(a)
    }

    dfs(graph, parent, 1, 0)

    for (i in 2..n) {
        println(parent[i])
    }
}

fun dfs(graph: Array<ArrayList<Int>>, parent: IntArray, node: Int, parentNode: Int) {
    parent[node] = parentNode 
    for (nextNode in graph[node]) {
        if (nextNode != parentNode) { 
            dfs(graph, parent, nextNode, node)
        }
    }
}