[백준] 집합의 표현 - 골드 4 (Kotlin)

문제

초기에 {0}, {1}, {2}, ... {n} 이 각각 n+1개의 집합을 이루고 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

입력

첫째 줄에 n(1 ≤ n ≤ 1,000,000), m(1 ≤ m ≤ 100,000)이 주어진다. m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다. a와 b는 n 이하의 자연수 또는 0이며 같을 수도 있다.

출력

1로 시작하는 입력에 대해서 한 줄에 하나씩 YES/NO로 결과를 출력한다. (yes/no 를 출력해도 된다)

소스 코드

import java.util.*

fun main() {
    val solution = Solution()

    solution.solution()
}

/**
 *
 */
class Solution {

    fun solution() {
        val (N, M) = readLine()!!.split(' ')
        val parent = Array(N.toInt() + 1) { 0 }

        for (i in parent.indices) {
            parent[i] = i
        }

        repeat(M.toInt()) {
            val input = readLine()!!.split(' ')

            when (input[0].toInt()) {
                0 -> {
                    unionSet(input[1].toInt(), input[2].toInt(), parent)
                }
                1 -> {
                    if (isUnion(input[1].toInt(), input[2].toInt(), parent))
                        println("YES")
                    else
                        println("NO")
                }
            }
        }
    }

    private fun findParent(n: Int, parent: Array<Int>): Int {
        if (parent[n] == n) return n
        parent[n] = findParent(parent[n], parent)
        return parent[n]
    }

    private fun unionSet(a: Int, b: Int, parent: Array<Int>) {
        val x = findParent(a, parent)
        val y = findParent(b, parent)

        if (x > y) parent[x] = y
        else parent[y] = x
    }

    private fun isUnion(a: Int, b: Int, parent: Array<Int>): Boolean {
        val x = findParent(a, parent)
        val y = findParent(b, parent)

        return x == y
    }
}

풀이 및 정리

Union Find 알고리즘을 처음 알게 되었다.
이 문제를 처음에는 내 방식대로 풀어보다가 메모리 초과, 시간 초과가 났고, 친구가 알고리즘을 알려줘서 쉽게 풀 수 있었다.
Union Find 는 합집합을 구할 때 사용할 수 있는 알고리즘이라고 했고 그 방식은 아래와 같다.
1. Parent 배열을 만들고 본인의 값으로 초기화한다. (배열, 트리 등 자료구조는 때에 따라 다른 것 같다)
2. 합집합을 만들 때는 두 집합의 원소의 parent 를 더 작은 쪽으로 (기준도 조건에 따라 다른 것 같다) 일치시킨다.
3. Find (집합의 부분 집합인지 판단) 할 때는 최상위 parent 가 일치하는지를 판단한다.
그림으로 설명하자면 (발로 그림)