백준 13164 - 행복 유치원

문제 링크: https://www.acmicpc.net/problem/13164

13164번: 행복 유치원

---

문제를 요약하면 다음과 같습니다.

• 주어진 오름차순으로 정렬된 배열을 \(A\)라 할 때, \(A\)를 \(K\)개의 조로 분할해야 합니다. 이때 \(K\)개의 조의 조원들은 서로 인접해 있어야 합니다.
• \(K\)개의 조로 나눌 때 아무렇게나 나누면 안되고, 각 조들의 비용 총합이 최소가 되도록 분할해야 합니다. 조의 비용은 해당 조의 (최댓값 - 최솟값)으로 계산합니다.

 

\(A\)가 정렬되어 있고, 각 조의 조원들이 서로 인접해 있으므로 각 조의 비용은 조원이 \(i, i + 1, i + 2, ... , j - 1, j\)일때 \(A_j - A_i\)가 됩니다.

---

이제 \(K\)개 조로 분할하는 방법을 생각해 봅시다. \(N\)명의 학생들을 \(K\)개 조로 분할해서 비용 총합이 최소인 것을 찾으려면 어떻게 해야 할까요?

가장 생각하기 쉬운 방법은 \(K\)개 조로 나누는 모든 방법을 시도하고, 이 중에 비용이 최소인 것을 고르는 것입니다. 이는 백트래킹을 사용하여 구할 수 있습니다.

하지만 이 방법은 \(N\)과 \(K\)의 제한이 \(300,000\)으로 높기 때문에 제한시간 내에 실행할 수 없습니다. 더 빠른 방법을 찾아야 합니다.

 

발상을 조금 바꿔봅시다. \(N\)명을 \(N\)개 조로 나눈 상황에서 시작해 봅시다. 이때 비용 총합은 당연하게도 \(0\)입니다.

여기서 \(N-1\)개 조로 조를 줄이려면, 기존에 있는 인접한 조 2개를 합쳐서 새로운 조를 하나 만들어야 합니다. 어떤 조 2개를 골라야 할까요?

원생 \(i\)의 키를 \(A_i\)라 하고, 원생 \(i, i+1, ..., j-1, j\)가 있는 조를 다음과 같이 표기하겠습니다.

\[[A_i, A_j]\]

이러면 \(N\)개 조를 다음과 같이 표현할 수 있습니다.

\[[A_1, A_1], [A_2, A_2], [A_3, A_3], ..., [A_N, A_N]\]

만약 조 \([A_1, A_1]\)과 \([A_2, A_2]\)를 합쳐 새로운 조 \([A_1, A_2]\)를 만들면 조의 개수는 \(N - 1\)개가 되고, 비용의 총합은 \(A_2-A_1\)이 됩니다.

그러면 \(A_{i+1}-A_i\)중에 가장 작은 것을 고른다면 조의 개수가 \(N - 1\)개 일때 비용 총합이 최소가 될 것입니다!

 

비슷한 방법으로 생각해 봅시다. 조가 \(T\)개 있다고 하고, \(T\)개 조가 다음과 같다고 합시다.

\[[A_1, A_{x_1}], [A_{x_1 + 1}, A_{x_2}], ..., [A_{x_{T-2}+1}, A_{x_{T-1}}], [A_{x_{T-1}+1}, A_N]\]

이 중 임의의 인접한 조 2개를 다음과 같이 골랐다고 해봅시다. 편의상 \(p, q, r\)의 표현을 사용합니다.

\[[A_p, A_q], [A_{q+1}, A_r]\]

이 두 조를 합쳐 새로운 조 \([A_p, A_r]\)를 만들게 되면 조의 개수는 \(T-1\)개가 되고, 비용의 총합이 변하게 됩니다.

기존 비용의 총합을 \(S_{old}\), 새로운 비용의 총합을 \(S_{new}\)라 하면 다음이 성립합니다.

\[S_{new}= S_{old} + (A_r - A_p) - (A_q - A_p) - (A_r - A_{q + 1}) = S_{old} + A_{q + 1} - A_q\]

즉 두 조를 합치면 두 조가 인접하는 부분에 있는 두 원소의 차이만큼 비용의 총합이 증가하게 됩니다! 

 

\(A\)는 오름차순 배열이므로 모든 \(q\)에 대하여 \(A_{q + 1} - A_q >= 0\)이 성립합니다. 따라서 비용 총합은 조가 합쳐져도 절대로 감소하지 않습니다.

따라서, 인접한 원소 차이가 가장 작은 것들을 계속 골라서 그 차이를 비용 총합에 더해주면 됩니다.

 

여기서 생각해 봐야할 것이 있습니다. 이전에 \(A_{x+1}\)와 \(A_x\)의 차이를 더해줬다면 또 이 값을 더할 수 있을까요?

\(A_x\)와 \(A_{x+1}\)이 더해졌다는 것은 이 두 값이 합쳐진 조의 양 끝이 아닌 중간에 들어갔다는 것을 의미합니다. 따라서 이 값들은 다시 사용할 수 없습니다.

즉, 어떤 \(x\)에 대하여 \(A_{x+1} - A_x\)는 비용 총합에 단 한 번만 더할 수 있습니다.

 

이제 문제를 해결할 수 있습니다. 인접한 두 원소의 차이 \(A_2 - A_1, A_3 - A_2, ..., A_N - A_{N-1}\)를 구하고, 오름차순으로 정렬합니다.

조를 한번 합칠 때 마다 조의 개수가 1개 줄어들고, 조가 \(N\)개인 상태로 시작하므로 총 \(N - K\)번 조를 합쳐주면 \(K\)개 조를 얻을 수 있습니다.

따라서 방금 정렬한 값들 중 앞의 \(N - K\)개를 더해주면, 비용 총합의 최솟값을 구할 수 있습니다.

시간 복잡도는 정렬과 덧셈을 합해 총 \(O(N \log N + N - K) = O(N \log N)\)입니다.

 

코드 (C++)

#include <bits/stdc++.h>

#define fastio ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);

using namespace std;
typedef long long ll;

int main() {
    fastio;
    
    int n, k;
    cin >> n >> k;
    
    ll cost = 0;
    vector<ll> height(n);
    
    for (int i = 0; i < n; i++) {
        cin >> height[i];
    }
    
    vector<ll> diff(n - 1);
    
    for (int i = 0; i < n - 1; i++) {
        diff[i] = height[i + 1] - height[i];
    }
    
    sort(diff.begin(), diff.end());
    
    for (int i = 0; i < n - k; i++) {
        cost += diff[i];
    }
    cout << cost;
    return 0;
}
// BOJ 13164