D - Various Sushi
概要
N個の寿司があり、i番目のネタはt[i], 美味しさはd[i]である。このうちK個食べたい。
食べたときの満足度は O + T となる。ただし
O = 食べた寿司の美味しさの総和
T = 食べたネタの種類の二乗
である。このとき、食べる寿司を上手く選んで満足度を最大化せよ。
感想
こういう問題は何かを固定して全探索をするのが定石で、この場合だと食べる種類数を全探索する、という方法はすぐに思いつきました。
しかし、ここで詰まります。種類数を固定しても「どのネタからいくつ取ってくる？」が決まらないからです。「あるネタから何個取ってくる」が決まれば、大きい美味しさ順に取ればいいのは明らかですが……どのネタをチョイスすべきか、これが分からない。
最適化問題なので満足度の二分探索なども考えるが、うまくいかなさそうです。DPはあるかもしれないですが、高々O(N)ぐらいまでしかできない、となると少し厳しそうでこの場合だと、どの種類までの寿司を考えたか、と、今まで何個食べたか、の情報ぐらいは必要そうだからO(N^2)はかかるんじゃないかっていうので却下しました。
となると、貪欲かなあという気持ちになりますが、どう貪欲するのかなあ……って感じです。
と、ここでたとえx種類の寿司を食べると仮定して実際はx種類以上食べていても大丈夫なんじゃないか？ということに気が付きます。
なぜならT = x^2なので過小評価することはあっても過大評価することはないからです。
というわけで次のような戦略を立てました。x = 1, 2, ...と動かしていって
として求めた値をmax更新していく。
実際には残りのK-x個の寿司を取るときに、最初に取ったx個以外のネタも混じっているかもしれないので種類数はxより大きくなっている可能性はあります。
ただそのときの解は、Tは過小評価してることになるし、和のほうもよく考えると過小評価していることになって、maxを取るときの影響を与えることがない、ということが言えます。（つまり最初に取ったx個以外のネタも混じっているような解は無視してよい）
この方法でやれば寿司をソートするのに最も時間がかかるのでO(NlogN)で解くことができます。
公式解説PDFの想定解と少し違うようなので久しぶりにブログ記事にしてみました。
