코틀린 코딩테스트 강좌, 최소 공통 조상 구하기 2

안녕하세요! 이번 시간에는 코틀린을 활용한 알고리즘 문제를 통해 최소 공통 조상(LCA, Lowest Common Ancestor) 구하기 문제를 해결해보겠습니다. 이 문제는 트리 구조를 다루는 중요한 개념 중 하나로, 다양한 코딩 테스트에서 등장합니다. 본 강좌는 주제에 대한 설명과 함께 문제 정의, 해결 과정, 그리고 최적화 방법까지 자세히 설명하겠습니다.

문제 정의

트리 구조가 주어졌을 때, 노드 A와 노드 B의 최소 공통 조상을 찾는 문제입니다. 최소 공통 조상은 노드 A와 노드 B의 조상 중에서 가장 깊은 조상을 의미합니다. 예를 들어, 아래와 같은 트리가 있다고 가정해보겠습니다.

이 경우, 노드 4와 노드 5의 최소 공통 조상은 노드 2입니다. 또한, 노드 4와 노드 6의 최소 공통 조상은 노드 1입니다.

입력

첫 번째 줄: 트리의 노드 개수 N (1 ≤ N ≤ 20,000)
두 번째 줄부터 N-1 번째 줄: 두 개의 정수 a, b가 주어짐. 이는 a가 b의 부모임을 의미.
마지막 줄: 두 개의 정수 X, Y (1 ≤ X, Y ≤ N) which are the nodes for which we need to find the LCA.

출력

최소 공통 조상의 노드 번호를 출력합니다.

예시 입력

예시 출력

문제 풀이 과정

이제 문제를 해결하기 위해 접근 방식을 설명하겠습니다. 트리 구조에서는 DFS(Depth-First Search)나 BFS(Breadth-First Search)를 활용하여 노드의 깊이를 저장하고 부모 노드를 기록하는 것이 일반적입니다. 이 방법을 통해 각 노드의 깊이를 비교하고, 공통 조상을 찾을 수 있습니다.

1. 트리 구축

우선, 입력으로 주어진 정보를 바탕으로 트리를 구축해야 합니다. 코틀린에서는 리스트를 사용하여 각 노드의 자식들을 저장하는 형태로 구현할 수 있습니다.


data class TreeNode(val id: Int) {
    val children = mutableListOf()
}

fun buildTree(n: Int, edges: List>): Map {
    val nodes = mutableMapOf()
    edges.forEach { (parentId, childId) ->
        val parentNode = nodes.getOrPut(parentId) { TreeNode(parentId) }
        val childNode = nodes.getOrPut(childId) { TreeNode(childId) }
        parentNode.children.add(childNode)
    }
    return nodes
}

2. 깊이 및 부모 노드 저장

DFS를 사용하여 깊이와 부모 노드를 저장하는 함수를 작성합니다. 이렇게 하면 나중에 LCA를 찾는 데 큰 도움이 됩니다.


val depth = mutableMapOf()
val parent = mutableMapOf()

fun dfs(node: TreeNode, dep: Int) {
    depth[node.id] = dep
    for (child in node.children) {
        parent[child.id] = node.id
        dfs(child, dep + 1)
    }
}

// 사용 예시
val root = nodes[1] // 루트 노드 (ID: 1)
dfs(root, 0)

3. LCA 찾기

이제 두 노드의 LCA를 찾는 함수를 구현합니다. 주어진 두 노드를 비교하여 깊이를 맞추고, 조상을 찾습니다.


fun findLCA(x: Int, y: Int): Int {
    // 깊이를 맞춤
    var a = x
    var b = y
    while (depth[a] > depth[b]) {
        a = parent[a]!!
    }
    while (depth[b] > depth[a]) {
        b = parent[b]!!
    }
    while (a != b) {
        a = parent[a]!!
        b = parent[b]!!
    }
    return a
}

// 사용 예시
val lca = findLCA(4, 6)
println("LCA: $lca")

4. 전체 구현

이제 모든 단계의 코드를 통합하여 전체 프로그램을 작성하겠습니다.


fun main() {
    val n = readLine()!!.toInt()
    val edges = mutableListOf>()
    repeat(n - 1) {
        val (a, b) = readLine()!!.split(" ").map { it.toInt() }
        edges.add(Pair(a, b))
    }
    val (x, y) = readLine()!!.split(" ").map { it.toInt() }

    val nodes = buildTree(n, edges)

    val depth = mutableMapOf()
    val parent = mutableMapOf()

    fun dfs(node: TreeNode, dep: Int) {
        depth[node.id] = dep
        for (child in node.children) {
            parent[child.id] = node.id
            dfs(child, dep + 1)
        }
    }

    val root = nodes[1]
    dfs(root, 0)

    fun findLCA(x: Int, y: Int): Int {
        var a = x
        var b = y
        while (depth[a] > depth[b]) {
            a = parent[a]!!
        }
        while (depth[b] > depth[a]) {
            b = parent[b]!!
        }
        while (a != b) {
            a = parent[a]!!
            b = parent[b]!!
        }
        return a
    }

    val lca = findLCA(x, y)
    println("LCA: $lca")
}

최적화

현재 방법은 DFS로 깊이와 부모 노드를 계산한 후 LCA를 구하는 방식입니다. 이 방식은 O(N)으로 동작하며, LCA를 단순히 찾는 데 O(log N)의 시간 복잡도를 가집니다. 하지만 보통 N이 커질 경우 최적화된 알고리즘인 “Sparse Table”이나 “Binary Lifting”을 활용할 수 있습니다. 이러한 방법들은 시간 복잡도를 더욱 줄일 수 있습니다.

결론

이번 강좌를 통해 최소 공통 조상을 찾는 방법에 대해 알아보았습니다. 트리 구조에서의 LCA 문제는 알고리즘 문제에서 자주 등장하는 주제이므로, 충분히 연습하여 몸에 익혀두는 것이 중요합니다. 각 단계별로 코드를 작성하고 이해하면서 연습해보세요. 다음 시간에는 다른 주제를 가지고 돌아오겠습니다. 감사합니다!