Codeforces Round #317 [AimFund Thanks-Round] (Div. 1) B. Minimization

Codeforces Round #317 [AimFund Thanks-Round] (Div. 1) (長い) に参加しました。
前考えていた通り B から解いてみたところ, B だけ正解してレートが上がりました。 A は解けなくてもレート上がるんや！

問題

codeforces.com

解法

まず K = 1 の場合を考えてみます。この場合は, 数列 A がソートされている時が最適なのは明らかでしょう。
では K が一般的な場合はどうなるかを考えてみます。この場合は,

A[0], A[k], A[2k], ...
A[1], A[k+1], A[2k+1], ...
...

がそれぞれソートされていれば良いです。
例えば, 問題文サンプルの 3 つ目のを考えてみます。

6 3
4 3 4 3 2 5

まず数列をソートすると

2 3 3 4 4 5

です。 K = 3 なので, 種類の区分は
A[0], A[3] と
A[1], A[4] と
A[2], A[5] です。今回は特に工夫するところもなく, 例えば
A[0] = 2, A[3] = 3
A[1] = 3, A[4] = 4
A[2] = 4, A[5] = 5
とソートされたものを当てはめれば最適解が得られます。

しかし, 例えば N = 7, K = 3 の場合, 種類の区分は
A[0], A[3], A[6] と
A[1], A[4] と
A[2], A[5] というように, (N%K) 個の区分けは N/K 個ではなく, N/K+1 個の区分けになります。よって, N/K 個の区分けをソートされた数列のどこに割り当てるか, N/K+1 個の区分けをソートされた数列のどこに割り当てるかが大事になるわけです。
そこで,

dp[k][rest] = (k 個の区分けをしているときに, N/K+1 個の区分けをすべき数列が残り rest 個になっている時の, 値の最小値)

という dp を作ると, k 個目の区分けは N/K 個にするか N/K+1 個にするかのどちらかしかないので, 各状態は O(1) で計算できます。よって, この dp を解くことによって答えを得ることが出来ます。

const int MAXN = 300010;
const int MAXK = 5005;
int A[MAXN];
ll memo[MAXN][2];
ll dp[MAXK][MAXK];
int N, K, R, Q;
const ll INF = 1ll<<61;

ll dfs(int k, int rest) {
    if (k == K) {
        if (rest > 0) return INF;
        return 0;
    }
    ll& ret = dp[k][rest];
    if (ret >= 0) return ret;
    ret = 0;
    ret = dfs(k+1, rest) + memo[k*Q+(R-rest)][0];
    if (rest > 0) {
        ret = min(ret, dfs(k+1, rest-1)+memo[k*Q+R-rest][1]);
    }
    return ret;
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cin >> N >> K;
    R = N%K;
    Q = N/K;
    for (int i = 0; i < N; i++) {
        cin >> A[i];
    }
    sort(A, A+N);
    {
        ll now = 0;
        for (int i = 0; i < Q-1; i++) now += A[i+1]-A[i];
        memo[0][0] = now;
        int cur = 0;
        for (; cur+Q < N; cur++) {
            now += (A[cur+Q]-A[cur+Q-1]);
            memo[cur][1] = now;
            now -= A[cur+1]-A[cur];
            memo[cur+1][0] = now;
        }
    }
    memset(dp, -1, sizeof(dp));
    cout << dfs(0, R) << endl;
    return 0;
}

面白い問題でした。