1717번: 집합의 표현 - Kotlin (Union - Find)

1717번: 집합의 표현

 

 

 

 

나의 풀이

이 문제는 Union - Find 알고리즘을 알고있다면 쉽게 해결할 수 있는 문제입니다. 만약 이 개념을 알지 못했다면 접근하기가 어려울 수 있습니다. Union - Find 알고리즘은 미리 정의된 원소들의 집합을 관리하고, 이들 집합 간의 합병(Union) 및 검색(Find)를 수행하는 알고리즘입니다.

 

스패닝 트리에 대해 알고 있다면 Union - Find 알고리즘을 사용해서 스패닝 트리를 구현할 수 있다는 것도 알고 있을 것입니다.

 

Union - Find 라는 이름에서 알 수 있듯이 해당 연산은 Union 연산과 Find 연산으로 이루어집니다.

 

 

Find 연산부터 알아보겠습니다.

fun find(x: Int): Int {
    return if (x != graph[x]) {
        graph[x] = find(graph[x])
        return graph[x]
    } else x
}

Find 연산은 특정 원소가 속한 집합을 찾는 연산입니다. 재귀함수를 통해 계속해서 그래프를 탐색하고 결국 루트 노드를 찾을 수 있습니다.

 

 

Union 연산은 Find 연산으로 찾은 두 원소의 루트 노드를 연결합니다.

fun union(x: Int, y: Int) {
    val tempX = find(x)
    val tempY = find(y)

    if (tempX != tempY) {
        graph[tempX] = tempY
    }
}

이런 식으로 루트 노드가 연결된다면 두 집합은 하나의 루트 노드를 가지게 됩니다.

 

 

그리고 만약 두 원소의 루트 노드가 일치한다면 둘은 같은 집합에 있다는 것을 의미합니다. 그렇기에 아래와 같이 함수로 집합 여부를 판단할 수 있습니다.

 

import java.util.*

lateinit var graph: Array<Int>

fun main() = with(System.`in`.bufferedReader()) {

    val (n, m) = readLine().split(" ").map { it.toInt() }

    graph = Array(n + 1) { it }

    repeat(m) {
        val (type, a, b) = readLine().split(" ").map { it.toInt() }

        when (type) {
            0 -> union(a, b)
            1 -> {
                if (find(a) == find(b)) {
                    println("YES")
                } else {
                    println("NO")
                }
            }
        }
    }
}

fun find(x: Int): Int {
    return if (x != graph[x]) {
        graph[x] = find(graph[x])
        return graph[x]
    } else x
}

fun union(x: Int, y: Int) {
    val tempX = find(x)
    val tempY = find(y)

    if (tempX != tempY) {
        graph[tempX] = tempY
    }
}