백준 17490 일감호에 다리 놓기 Kotlin (크루스칼)

문제 출처 : https://www.acmicpc.net/problem/17490

17490번: 일감호에 다리 놓기

문제

학교의 홍보대사를 맡게 된 건덕이는 건국대학교의 모든 강의동을 신입생들에게 소개해야 한다.

건국대학교 중앙에 위치한 일감호를 따라 한 바퀴를 돌며 모든 강의동을 소개하는 것이 그의 일이지만, 몇몇 구간들이 공사 중이어서 그 구간을 통해서는 갈 수 없는 상황이다. 급한대로 건덕이는 호수에 돌을 던져 징검다리를 놓아 길을 만들어보려고 한다.

강의동은 일감호의 둘레에 따라 원형으로 배치돼 있으며, 강의동 양 옆의 강의동은 서로 이웃한다. 또, 원형으로 배치돼 있기 때문에 N개의 강의동이 있다면 N번째 강의동과 1번째 강의동은 서로 이웃한다.

일감호 안에는 와우도라는 섬이 있다. 건덕이는 한 강의실에서 다른 모든 강의실로 이동할 수 있도록 강의동에서 와우도까지 징검다리를 놓기로 했다. 하지만 건덕이의 눈에는 K개의 돌밖에 보이지 않는다. 건덕이는 주어진 돌을 활용해서 징검다리를 완성할 수 있을까?

입력

첫째 줄에 강의동의 수 N, 공사구간의 수 M, 건덕이가 가진 돌의 수 K가 공백으로 구분돼 주어진다. 강의동은 1동부터 N동까지 존재한다.

다음 줄에는 강의동에서 와우도까지 놓아야하는 돌의 개수 S1, S2, ..., SN이 공백으로 구분돼 주어진다. 이는 T번째 강의동에서 와우도까지 ST개의 돌을 놓아야 함을 의미한다. 이어서 M개의 줄에 i, j가 주어진다. 이는 i번째 강의동에서 j번째 강의동까지 가는 길이 공사중임을 의미한다. 이 때 입력되는 i, j번째 건물은 이웃한 강의동이다. 공사중인 구역은 한 번만 주어진다.

출력

건덕이가 가지고 있는 돌을 놓아 모든 강의동을 연결할 수 있으면 YES를, 그렇지 않으면 NO를 출력한다.

제한

• 3 ≤ N ≤ 1,000,000
• 0 ≤ M ≤ N
• 1 ≤ i, j ≤ N 
• 1 ≤ T ≤ N인 모든 정수 T에 대해 1 ≤ ST ≤ 1,000,000
• 0 ≤ K ≤ 5,000,000,000

알고리즘 분류

• 그래프 이론
• 자료 구조
• 분리 집합
• 최소 스패닝 트리

풀이

최소 스패닝 트리를 구하는 문제이다.

사실 이 문제는 그리디로도 풀 수 있지만  패쓰

본인은 크루스칼과 프림 중 크루스칼로 풀었다.

우선 공사중인 길이 없다고 하면 5개의 강의동이 있다고 할 때,

1-2

2-3

3-4

4-5

5-1

모두 연결되어있기 때문에 이를 간선으로 표현할 수 있고

parent 배열이 다음과 같이 있다고 하면 여기에 이 간선들로 UnionParent함수로 그룹화할 수 있다.

before

index	0	1	2	3	4	5
value	0	1	2	3	4	5

after

index	0	1	2	3	4	5
value	0	1	1	1	1	1

이제 공사된 길을 표현해 보자.

공사된 길의 입력은 꼭 i, i+1의 구조로 나온다는 보장이 없음에 유의하자.

만약 2-3의 길이 공사중이라고 하면 3-2라고 나올 수도 있다.

따라서 3-2라고 나오든 2-3으로 나오든 작은 값을 from, 큰 값은 to로 하면

destroyed[to]로 공사로 인해 끊어진 길을 표현할 수 있다.

그럼 UnionParent로 간선을 연결할 때 destroyed[to]가 true인 간선은 연결 스킵하면 된다.

그럼 1번 예제에서 유효한 간선은

1-2

3-4

두 개 뿐이다.

그럼 이 간선들만 UnionParent하면

index	0	1	2	3	4	5
value	0	1	1	3	3	5

이렇게 {1,2}, {3,4}, {5} 3개의 그룹으로 나누어진다.

그럼 이제 섬에서 이 세 그룹에 징검다리를 연결하면 되는데 각각의 그룹에서 가장 비용이 적은 다리를 연결하면 된다.

 

그럼 각각 그룹에서 비용이 적은 다리를 어떻게 알 수 있을까?

애초에 UnionParent로 그룹화할 때 비용이 적은 강의동으로 합쳐주면 된다.

예를 들어 {1,2} 중에 2가 다리를 연결하는 데 비용이 더 적고, {3,4}중에 4가 더 적다면

index	0	1	2	3	4	5
value	0	2	2	4	4	5

이렇게 된다.

이제 섬에서 2,4,5에 연결하는 다리 비용들을 더해주고 k값과 비교하면 된다.

 

주의할 점은 K가 Int의 범위를 넘는다는 점, 정답도 Int 범위를 넘을 수 있다는 점이다.

코드

import java.util.*
val br = System.`in`.bufferedReader()
fun getIntList() = br.readLine().split(' ').map { it.toInt() }
fun getInt() = br.readLine().toInt()

lateinit var destroyed: BooleanArray
lateinit var parent: IntArray
lateinit var costs: List<Int>
fun getParent(x: Int): Int {
    return if (x != parent[x]) getParent(parent[x]).also { parent[x] = it } else x
}

fun unionParent(x: Int, y: Int) {
    val xx = getParent(x)
    val yy = getParent(y)
    if (costs[xx] > costs[yy]) {
        parent[xx] = yy
    } else {
        parent[yy] = xx
    }
}

fun main() = with(System.out.bufferedWriter()) {
    //input
    var tk = StringTokenizer(br.readLine())
    val n = tk.nextToken().toInt()
    val m = tk.nextToken().toInt()
    val k = tk.nextToken().toLong()
    parent = IntArray(n + 1) { it }
    destroyed = BooleanArray(n + 1)
    tk = StringTokenizer(br.readLine())
    costs = List(n + 1) { if (it == 0) 0 else tk.nextToken().toInt() }

    for(i in 0 until m) {
        tk = StringTokenizer(br.readLine())
        val (a, b) = List(2){tk.nextToken().toInt()}
        val large = a.coerceAtLeast(b)
        val small = a.coerceAtMost(b)
        if (large == n && small == 1) {
            destroyed[1] = true
        } else {
            destroyed[large] = true
        }
    }
    //solve
    //makeUnion
    for (i in 1..n) {
        var j = (i + 1) % (n + 1)
        if (j == 0) j = 1
        if(destroyed[j]) continue
        unionParent(i,j)
    }
    //check
    var answer = 0L
    var unionCnt=0
    for(i in 1 .. n){
        if(getParent(i) == i){
            answer += costs[i]
            unionCnt++
        }
    }
    //output
    write(if (answer <= k || unionCnt <= 1) "YES" else "NO")
    close()
}