Yuki 615 集合に分けよう

Problem Description:
There is a set of N numbers. You would like to split the set into K non-empty subsets, so that the sum of contribution of each subset is minimal. Contribution of a subset can be computed as the absolute difference between the maximum element and the minimum element in the subset.

Constraints:
1 ≤ K ≤ N ≤ 5e4
0 ≤ A[i] ≤ 1e12

Solution:
Consider how the answer changes when K increases. In the beginning, K = 1, the answer is max(A) - min(A). Then let K = 2, we could observe that we can remove the greatest difference between any two adjacent values, in the sorted array, from the answer at K = 0. This also works for any positive integer for K, hence we would subtract the greatest K - 1 difference of adjacent values in the sorted array from max(A) - min(A), leading to the optimal answer.

Code:

#include <bits/stdc++.h>

long long greedy(int n, int k, const std::vector<long long> &a) {
  std::vector<long long> aa = a;
  std::sort(aa.begin(), aa.end());
  std::vector<long long> xa(n - 1);
  for (int i = 0; i + 1 < n; ++i) {
    xa[i] = aa[i + 1] - aa[i];
  }
  std::sort(xa.begin(), xa.end(), std::greater<long long>());
  long long res = aa.back() - aa.front();
  for (int i = 0; i + 1 < k; ++i) {
    res -= xa[i];
  }
  return res;
}

signed main() {
  std::ios::sync_with_stdio(false);
  int N, K;
  std::cin >> N >> K;
  std::vector<long long> A(N);
  for (int i = 0; i < N; ++i) {
    std::cin >> A[i];
  }
  std::cout << greedy(N, K, A) << std::endl;
  return 0;
}