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; }