[백준] 1717번: 집합의 표현 - Kotlin[코틀린]

문제

 

1717번: 집합의 표현

 

---

풀이

 

유니온 파인드(Union-Find)를 활용한 문제다. 0으로 시작하는 입력은 두 집합을 합치는 합집합(Union) 연산을 수행하고, 1로 시작하는 입력은 두 원소가 같은 집합에 속하는지 확인한다.

 

 유니온 파인드란 각 노드가 어떤 집합에 속하는지 판별하는 알고리즘이다. 이를 배열을 이용해 구현하는데 초기에는 배열의 각 요소에는 자신을 가리키도록 한다. Union 연산을 수행할 때는 배열의 값을 루트 노드의 값으로 변경한다. 일반적으로 낮은 인덱스을 루트 노드로 설정한다. 예를 들어 배열의 초기 상태가 {1, 2, 3, 4, 5}라면 Union연산으로 1번과 2번을 합치면 배열은 {1, 1, 3, 4, 5}로 변경된다. 그리고 2번과 3번을 합치면 배열은 {1, 1, 1, 4, 5}가 된다. 이러한 상태는 1, 2, 3은 합집합이고 나머지는 아닌 경우가 되는 것이다.

 

---

코드

 

import java.util.*

fun main() {
    val br = System.`in`.bufferedReader()
    val bw = System.out.bufferedWriter()

    val sb = StringBuilder()

    val st = StringTokenizer(br.readLine())
    val num = st.nextToken().toInt()
    val case = st.nextToken().toInt()

    val parent = IntArray(num+1) { it }

    fun find(a: Int): Int { // 루트 노드를 찾는 연산
        return if(a == parent[a]) a // 배열의 인덱스와 값이 같음
        else { // 부모 노드의 번호를 전달하면서 루트 노드를 찾을 때까지 재귀 호출
            parent[a] = find(parent[a])
            parent[a]
        }
    }
    fun union(a: Int, b: Int) { // 두 노드를 같은 집합으로 합치는 연산
        val x = find(a)
        val y = find(b)
        if(x<y) { // 작은 값이 부모 노드가 되도록 합침
            parent[y] = x
        } else {
            parent[x] = y
        }
    }

    repeat(case) {
        with(StringTokenizer(br.readLine())) {
            if(nextToken().toInt() == 0) { // 0인 경우 유니온 연산
                union(nextToken().toInt(), nextToken().toInt())
            } else { // 1인 경우 같은 집합인지 확인
                sb.append(if (find(nextToken().toInt())==find(nextToken().toInt())) "YES\n" else "NO\n")
            }
        }
    }

    bw.write("$sb")
    bw.flush()
    bw.close()
    br.close()
}