백준 13164 행복 유치원

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

N명의 원생들을 K개의 그룹으로 나누는데, 그룹마다 티셔츠 비용 (그룹 내 최장신 키 - 최단신 키)이 들고 모든 비용 합이 최소일 때 비용을 출력하는 문제이다.

원생의 수는 최대 300,000이고 조의 개수도 최대 N과 동일하므로 O(NlogN) 이하로 처리해야 하는 문제이다.

예제를 통해 떠오른 접근 아이디어는 아래와 같다.

모든 원생들을 그룹으로 나눌 때, 비용이 가장 적게 드는 경우는? 모두 따로 조를 구성했을 때 비용이 0원이다.

[1, 3, 5, 6, 10] 의 입력이 주어졌다.

1 | 3 | 5 | 6 | 10 으로 나누는 경우가 가장 비용이 적다.

하지만 이 상태에서 그룹을 K개 (예시 기준 3개)로 만들어야 한다.

그러면 그룹을 나누는 4개의 막대기( | )들 중 2개 (5 - 3)를 치우는 문제가 된다.

여기서 비용이 가장 적게 들도록 하려면 제거할 막대기 양 옆의 키 차이가 가장 적은 막대기부터 치우면 된다.

1 | 3 | 5 | 6 | 10 에서 각 차이는 2, 2, 1, 4이다.

따라서 차이가 가장 적은 5와 6 사이의 막대기를 먼저 치운다.

그리고 다음으로 차이가 2인 막대기를 치우는데, 편의상 앞의 것부터 치우도록 한다.

키 차이가 최소인 막대 2개를 치움으로써 비용이 가장 적게 드는 3개의 그룹으로 원생들을 나누게 된다.

fun main(args: Array<String>): Unit = with(System.`in`.bufferedReader()) {
    val (N, K) = readLine().split(' ').map{it.toInt()}
    val students = readLine().split(' ').map{it.toInt()}

    // 키 차이를 각각 구해 정렬한 것 (i번째 차이 = i + 1번째 학생 키 - i번째 학생 키)
    val gaps = Array(N - 1){
        Gap(it, it + 1, students[it + 1] - students[it])
    }.sortedArrayWith{g1, g2 -> g1.gap - g2.gap}

    // i번째 학생과 i + 1번째 학생 사이의 막대기 유무
    val splits = BooleanArray(N - 1){ true }

    // 양 옆 키 차이가 적은 막대기부터 N - K개 제거 (K개 그룹 만들기 위함)
    repeat(N - K){
        val next = gaps[it]
        splits[next.idx1] = false
    }

    var answer = 0
    var start = 0
    // 학생 배열 순회하며 막대기 없으면, 첫 학생과 마지막 학생 키 차이만큼 answer에 더함
    for(i in students.indices){
        if(start == 0 && i < splits.size && splits[i]) continue
        if(start == 0)
            start = students[i]
        if(i == splits.size || splits[i]){
            answer += students[i] - start
            start = 0
        }
    }

    print(answer)
}

data class Gap(
    val idx1: Int,
    val idx2: Int,
    val gap: Int
)