[백준 10838번] [Kotlin] 트리

 

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

10838번: 트리

 

난이도 : 플래티넘 5
태그 : 트리, 최소 공통 조상

 

 

설명

최소 공통 조상(LCA)를 사용하는 문제이지만, 일반적인 LCA와 달리, 새로운 사용법을 알게 된 문제였습니다.

 

각 노드의 깊이를 찾기 위해

일반적으로 루트 노드에서 DFS 탐색을 진행하여 각 노드의 부모와 깊이를 알아냅니다.

 

LCA에서 depth를 구하는 이유는 최소 공통 조상을 찾기 위한 과정 중,

v1, v2 두 노드가 주어졌을 때, 둘 중 깊이가 낮은 노드를 기준으로 depth를 같게 맞춰준 뒤,

한 단계식 부모 노드로 올라가며 같은 노드가 나올때 까지 반복하기 때문입니다.

 

하지만 해당 문제의 경우, 노드의 부모를 계속 바꾸기 때문에,

여러번의 DFS를 돌려야 하고, 때문에 depth를 이용해 풀게 되면 시간초과가 발생할 것 같아 다른 방법을 찾아야 할 것 같습니다.

 

 

 

 

주목해야 할 부분은 위 문구입니다.

 

paint와 count 연산 시 두 노드의 최단경로는 1000이하 이므로,

v1, v2의 부모를 타고 올라가, 적어도 각각의 1000번째 부모 안에 최소 공통 조상이 있음을 의미합니다.

 

즉, v1의 부모를 1000번째 까지 타고 올라가며 방문체크를 해준 뒤,

v2의 부모를 1000번째 까지 타고 올라가면서, 이미 방문한 노드가 있다면?

해당 노드가 최소 공통 조상 노드가 되는 것 입니다.

 

1000이하이기에, 브루트포스적으로 접근하여도 시간안에 풀 수 있습니다.

 

 

 

 

 

 

소스코드

import java.util.StringTokenizer

data class Node(var parent: Int, var color: Int)

var n = 0
lateinit var nodes: Array<Node>

fun main() = with(System.`in`.bufferedReader()) {
    val nk = readLine().split(" ").map { it.toInt() }
    n = nk[0]
    val k = nk[1]

    nodes = Array(n) { Node(0, 0) }
    nodes[0].parent = -1

    val sb = StringBuilder()

    repeat(k) {
        val st = StringTokenizer(readLine())
        val (cmd, v1, v2) = Array(3) { st.nextToken().toInt() }

        when (cmd) {
            1 -> {
                val color = st.nextToken().toInt()
                paint(v1, v2, color)
            }
            2 -> move(v1, v2)
            3 -> sb.appendLine(count(v1, v2))
        }
    }

    print(sb)
}

fun LCA(v1: Int, v2: Int): Int {
    val visited = BooleanArray(n)
    var cur = v1

    for (i in 0 until 1000) {
        visited[cur] = true
        cur = nodes[cur].parent

        if (cur == -1) break
    }

    cur = v2
    for (i in 0 until 1000) {
        if (visited[cur]) return cur
        cur = nodes[cur].parent

        if (cur == -1) break
    }

    return -1
}

fun paint(v1: Int, v2: Int, color: Int) {
    val foundLCA = LCA(v1, v2)

    if (foundLCA == -1) return

    var cur = v1
    while (cur != foundLCA) {
        nodes[cur].color = color
        cur = nodes[cur].parent
    }

    cur = v2
    while (cur != foundLCA) {
        nodes[cur].color = color
        cur = nodes[cur].parent
    }
}


fun count(v1: Int, v2: Int): Int {
    val foundLCA = LCA(v1, v2)
    val colorSet = HashSet<Int>()

    var cur = v1
    while (cur != foundLCA) {
        colorSet.add(nodes[cur].color)
        cur = nodes[cur].parent
    }

    cur = v2
    while (cur != foundLCA) {
        colorSet.add(nodes[cur].color)
        cur = nodes[cur].parent
    }

    return colorSet.size
}

fun move(v1: Int, v2: Int) {
    nodes[v1].parent = v2
}

 

 

전체 소스코드입니다.

 

main 함수의 경우, 인풋값 처리 및 저장, cmd 값에 따른 메소드 호출이기 때문에,

주목할 부분은 LCA, count, move, paint 메소드입니다.

 

 

 

 

move(v1, v2)

fun move(v1: Int, v2: Int) {
    nodes[v1].parent = v2
}

 

move는 v1의 부모를 v2로 만들어주는 메소드입니다.

 

 

 

 

 

LCA(v1, v2)

fun LCA(v1: Int, v2: Int): Int {
    val visited = BooleanArray(n)
    var cur = v1

    for (i in 0 until 1000) {
        visited[cur] = true
        cur = nodes[cur].parent

        if (cur == -1) break
    }

    cur = v2
    for (i in 0 until 1000) {
        if (visited[cur]) return cur
        cur = nodes[cur].parent

        if (cur == -1) break
    }

    return -1
}

 

LCA는 최소 공통 조상(LCA)를 구하는 메소드입니다.

 

앞서 언급했듯이, depth를 쓰지 않고,

v1의 1000번째 부모까지 방문체크를 하며 올라가고,

v2의 1000번째 부모까지 올라가되, 이미 방문한 노드를 찾았다면,

 

해당 노드가 v1, v2의 최소 공통 조상이 됩니다.

 

 

 

 

 

paint(v1, v2)

fun paint(v1: Int, v2: Int, color: Int) {
    val foundLCA = LCA(v1, v2)

    if (foundLCA == -1) return

    var cur = v1
    while (cur != foundLCA) {
        nodes[cur].color = color
        cur = nodes[cur].parent
    }

    cur = v2
    while (cur != foundLCA) {
        nodes[cur].color = color
        cur = nodes[cur].parent
    }
}

 

paint의 경우, lca를 구한 뒤,

v1 ~ lca, v2 ~ lca 구간에 색상을 칠하는 메소드입니다.

 

 

 

 

 

count(v1, v2)

fun count(v1: Int, v2: Int): Int {
    val foundLCA = LCA(v1, v2)
    val colorSet = HashSet<Int>()

    var cur = v1
    while (cur != foundLCA) {
        colorSet.add(nodes[cur].color)
        cur = nodes[cur].parent
    }

    cur = v2
    while (cur != foundLCA) {
        colorSet.add(nodes[cur].color)
        cur = nodes[cur].parent
    }

    return colorSet.size
}

 

count 메소드는, lca를 구한 뒤,

v1 ~ lca, v2 ~ lca 구간에서 나온 서로 다른 색상의 개수를 출력하는 메소드입니다.

서로 다른 색상의 개수를 구하는 것이기 때문에, 저는 Set을 사용하여 중복 제거를 해줬습니다.

 

 

 

 

 

후기

LCA... 이긴 한데 이런 문제는 처음 봅니다.

두 노드의 최단 거리 제한을 두어 depth를 사용하지 않고 최소 공통 조상을 찾는, 재밌었던 문제였던 것 같네요.