[백준] 1976번: 여행 가자 - Kotlin[코틀린]

문제

 

1976번: 여행 가자

 

---

풀이

 

 유니온 파인드(Union-Find)를 활용한 문제다. 도시들의 연결 정보를 입력받아 합집합(Union) 연산을 수행하고 여행 계획에 포함되어 있는 도시들이 합집합인지 아닌지 판별하여 결과를 출력하면 된다.

 

---

코드

 

import java.util.*

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

    var st : StringTokenizer
    val num = br.readLine().toInt()
    val city = br.readLine().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
        }
    }

    for(i in 1 .. num) {
        st = StringTokenizer(br.readLine())
        for(j in 1 .. num) {
            if(st.nextToken().toInt() == 1) { // 1인 경우 유니온 연산
                union(i, j)
            }
        }
    }
    st = StringTokenizer(br.readLine())
    val start = parent[st.nextToken().toInt()]
    var check = "YES"
    for(i in 1 until city) {
        if(parent[st.nextToken().toInt()]!=start) { // 도시가 합집합이 아닌 경우
            check = "NO"
            break
        }
    }

    bw.write(check)
    bw.flush()
    bw.close()
    br.close()
}