[백준 16562번] [Kotlin] 친구비

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

16562번: 친구비

 

난이도 : 골드  4
태그 : 자료구조, 그래프이론, 그래프탐색, 분리집합

 

 

설명

친구의 친구는 나에게도 친구이므로,

한 친구만 만들면 해당 친구가 속한 그룹은 모두 내 친구가 되므로,

그룹 당 친구비가 가장 적은 친구와 친구를 맺으면 됩니다.

 

 

 

접근방법

1. 분리집합 (유니온 - 파인드)를 사용해 각 친구관계를 그룹화한다.

2. 각 친구 그룹의 최상위 노드를 Set에 저장한다.

3. 친구비가 가장 적은 순으로 해당 친구의 root를 Set에서 제거하고, 총 비용을 누적한다.

4. Set이 빌 때 까지 반복한다.

 

 

먼저, 친구의 친구는 자동으로 친구가 되므로,

각 친구관계를 그룹화하는 과정이 필요합니다

해당 과정은 저는 분리집합 / Union-Find를 사용하였습니다

 

색칠된 노드는 각 그룹의 루트 노드입니다.

 

 

 

그리고, 각 그룹의 루트 노드만 Set에 저장합니다.

 

 

 

(친구번호, 친구비)를 친구비 기준으로 오름차순 정렬합니다.

 

 

1번 친구와 친구를 맺겠습니다.

1번 친구의 root는 3번으로, Set에서 제거합니다.

Set에 남은 원소는 4이며, 현재까지 소비한 비용은 10입니다.

 

그 다음으로 친구비가 낮은 2번 친구와 친구를 맺습니다.

2번 친구의 root 노드는 4번 노드입니다.

Set에서 4번 노드를 꺼내고, Set이 비었으므로 프로그램을 종료합니다.

 

총 비용은 20 입니다. 

 

 

 

소스코드

import java.util.*
import kotlin.collections.HashSet

data class Cost(val index: Int, val cost: Int)

lateinit var parent: IntArray

fun main() = with(System.`in`.bufferedReader()) {
    val (n, m, k) = readLine().split(" ").map { it.toInt() }
    val st = StringTokenizer("0 ${readLine()}")
    val costs = Array(n + 1) { Cost(it, st.nextToken().toInt()) }.sortedBy { it.cost }

    parent = IntArray(n + 1) { it }

    for (i in 0..<m) {
        val (a, b) = readLine().split(" ").map { it.toInt() }
        union(a, b)
    }

    val friendsRoot = HashSet(parent.map { find(it) })

    var ans = 0

    for (i in 1..n) {
        val root = find(costs[i].index)
        if (root in friendsRoot) {
            friendsRoot.remove(root)
            ans += costs[i].cost
        }

        if (friendsRoot.isEmpty()) break
    }

    println(if (ans <= k) ans else "Oh no")
}

fun find(x: Int): Int {
    if (parent[x] == x) return x
    parent[x] = find(parent[x])
    return parent[x]
}

fun union(x: Int, y: Int) {
    val (px, py) = find(x) to find(y)
    parent[py] = px
}

 

 

 

후기

준석이에게 친구가 많이 생겼으면 좋겠네요.