[Kotlin] 백준 1717 집합의 표현

📘 문제

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

💡 풀이

union-find 알고리즘을 공부하던 중 발견한 문제여서 접근은 처음부터 union-find를 쓰면 되겠구나라고 생각했다.
m개의 줄에는 각각 0,a,b or 1,a,b의 형태로 연산이 주어진다.
이 때 0은 합집합, 1은 a와 b가 같은 집합에 포함되어 있는지를 확인하여 "YES"나 "NO"를 출력하면된다.

합집합일 때는 union으로 두 개의 원소를 같은 부모를 가리키게 만든다. 즉, 같은 집합에 속하게끔.
1일 때는 두 개의 원소를 find하여 같은 부모를 가리키고 있다면, 같은 집합에 속해있는 것으로 판단하여 "YES"를 출력하면 된다.

위에서 말하는 '같은 부모'를 비교할 때는 반드시 '조상'(연결된 노드 중 가장 위의 노드)으로 비교해야 한다.

💻 코드

private lateinit var nums: IntArray
private fun main() {
    val (n, m) = readln().split(" ").map { it.toInt() }
    nums = IntArray(n + 1) { it }
    for (i in 0 until m) {
        // 0: union, find: 1 a b
        val (type, a, b) = readln().split(" ").map { it.toInt() }
        if (type == 0) {
            union(a,b)
        } else {
            if (isContains(a, b)) println("YES")
            else println("NO")
        }
    }
}
private fun find(x: Int): Int {
    if (nums[x] == x) return x
    nums[x] = find(nums[x])
    return nums[x]
}

private fun union(x: Int, y: Int) {
    val findX = find(x)
    val findY = find(y)
    if (findX == findY) return
    if (nums[findY] < nums[findX]) nums[findX] = findY
    else nums[findY] = findX
}

private fun isContains(x: Int, y: Int): Boolean = find(x) == find(y)

결과