CODE FESTIVAL 2018 Final: G - Chicks and Cages

https://atcoder.jp/contests/code-festival-2018-final-open/tasks/code_festival_2018_final_g

問題概要

N羽のひよこがいて、ひよこiの大きさはAiである。ひよこをM個のケージに分けて入れるとき、各ひよこは自分とおなじケージにいるひよこの大きさの和と等しいストレスを受ける。全ひよこのストレスの和を最小化するようにひよこを分けたとき、その最小値を求めよ。

N, M <= 2000

解法

DPをしようと思うと、以下のような形になる。

DP(i, j): i匹目までのひよこをj個のケージに分けて入れたときの最小値

このDPで遷移を普通にやろうとすると次のグループの大きさを全て見る必要があるため、累積和を使ってもO(N)かかってしまい、状態数のO(NM)と合わせてO(N2 M)と厳しいオーダーになってしまう。

想定解ではこの遷移を高速化する。そのために必要なアイデアは「数の多いグループには小さいひよこを入れた方が良い」ということである。ひよこiが羽数nのケージに入った場合、最終的な値には A[i] × n が加算される。つまり全部のひよこについての A[i] × (ケージの羽数) を足した数が最終的な値ということになる。よって仮にすべてのケージの羽数が定まっているならば、A[i]の小さいひよこをnの大きいケージに貪欲に割り当てていくのが当然最適となる。つまりAを小さい順に並べて、大きいケージから順番に詰め込んでいくことになる。

以上の考察より、DPの遷移を次のとおり高速化できる。まずAを小さい順にソートする。先の考察に従えば、1つのグループはソート後のAの連続部分列になる。さらに「小さいひよこほど大きいケージへ」が最適であるため、後ろのひよこに作るケージは前に作るケージと同じサイズかもしくは小さいはずである。これによって考慮するべきケージのサイズが後半ほど枝刈りできて通るようになる(前からケージを作っていくとして、これからサイズnのケージを作るとき、これまで作ったm個のケージはすべてサイズn以上と仮定できる。ということはこれまでの割当でひよこを最低でも n * m 匹は使っているはずで、もしこの数が今見てるひよこのインデックスを追い越していたらNGなので、そのサイズ以上のケージの計算は飛ばせる)。

公式解説によると上の方法をとった場合の計算量はO(NMlogN).

感想

いや、わからん……………… 「ひよこiの寄与はA[i]*ケージのサイズ」に着目しないと始まらない感じだが、着目できていたとてDPのこういう高速化に使う発想はちょっと出てきそうにない うぐぐ

コード

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, core.stdc.string;

void main() {
    auto s = readln.split.map!(to!int);
    auto N = s[0];
    auto M = s[1];
    auto A = readln.split.map!(to!long).array;
    A.sort();

    auto B = new long[](N+1);
    foreach (i; 0..N) B[i+1] = B[i] + A[i];

    auto dp = new long[][](N+1, M+1);
    foreach (i; 0..N+1) fill(dp[i], 1L<<59);
    dp[0][0] = 0;

    foreach (i; 0..N) {
        foreach (j; 0..M) {
            foreach (k; i..N) {
                int size = k - i + 1;
                if (size * j > N) break;
                dp[k+1][j+1] = min(dp[k+1][j+1], dp[i][j] + (B[k+1] - B[i]) * size);
            }
        }
    }

    dp[N][M].writeln;
}