ABC116 D - Various Sushi - ツバサの備忘録

解法

ある程度までは貪欲にとっていき、その後一部入れ替えてみて探索していく、という形になります。
まずは、全ての寿司を美味しさの降順でソートします。
そして、上位 $K$ 個を食べたときのスコアを計算します。
この時、種類数を数えるのと同時に、種類 $i$ の寿司をいくつ食べたか、も記録しておきます。
そして、 $c$ 種類で美味しさの総和が $p$ であったとします。
このときのスコア、 $p + c^{2}$ は解の候補の1つになります。
以降は、 $K+1$ 個目以降の寿司についてみていきます。今 $j$ 番目の寿司を見ているとします。
もし、 $j$ 番目の寿司の種類が、今食べることにしている $K$ 個の寿司の種類と被っていない場合、以下のような操作をします。

今食べることにしている寿司のうち、同じ種類で2個以上存在する寿司の中から、最もおいしさが少ないものと交換する。

結局やりたいことは、種類ボーナスの獲得をするために、種類が増えるようにかつ、一番美味しさの総和が減らないようにしながら食べる寿司を交換する、ということです。
これを $K+1$ 番目以降のすべての寿司について行います。
具体的な実装方法としては、

$j$ 番目の寿司の種類を確認し、今食べることにしている $K$ 個の寿司と種類が被っていないか確認する(被っていたら $j+1$ へ進む)
$r$ (初期値は $K$ で、 $j+1$ 以降も使いまわします)番目のお寿司から順番に見ていき、 $r$ 番目以外にも、同じ種類のお寿司が $K$ 個の中にあれば、 $r$ 番目のお寿司を食べないようにし、 $j$ 番目のお寿司を食べることにする。他に存在しなければ、 $r$ の番号を若くしていき、同じことを繰り返す(もしすべて1個ずつしかなければ、その時点で全ての操作は終了します)。

とすると、おいしさの総和の減少を最小限に抑えつつ、種類のみを増やすことができます。あとは、これを繰り返し、食べる寿司を交換するたびに最大値をとっていけば答えが求まります。

では、なぜこのような貪欲法で通るのでしょうか。
まず、最初においしさの上位 $K$ 個の寿司を食べると仮定したので、その種類数 $c$ 以下になるような食べ方では、絶対にスコアが最大になることはありません。美味しさの総和でも、種類ボーナスでも最初の値を超えることができないからです。ということで、 $K+1$ 番目以降のお寿司を採用する際は、種類が増えるように選択しなければなりません。よって、 $K+1$ 番目以降のお寿司で、上位 $K$ 個のいずれかと種類が同じ寿司を採用すること、そして $K$ 個の寿司の中で、1個しか存在していない種類の採用をとりやめることはありえなくなります。
さて、あとは $K+1$ 以上の2つの数字 $j,l (j \lt l)$ について、おいしさが $j$ 番目の寿司を採用せずに $l$ 番目の寿司のみを採用するパターンがあるかどうか、です。どちらの場合も、その寿司を採用したことで増える寿司の種類は1種類です。よって、種類ボーナスの上昇量は等しいため、よりおいしさが高い $j$ 番目の寿司を採用する方が、スコアは当然高くなります。
そのため、 $j,l$ 両方を採用するパターンはあれど、 $l$ のみを採用するパターンはありえません。
ということで、 $K+1$ 番目からは、順番に種類が増えるようにお寿司を採用(交換)していけば、最大値を探索することができます。

感想

この手の問題はとっつきにくい気がします。苦手意識があるのですが、わりとはやいスピードで解くことができたのでよかったです。
このような問題では、とりあえずDPか貪欲が自分の中では思い浮かぶのですが、評価値が複雑なため、何か貪欲できる部分がありそう、という考えに至ります。そこで、まずは種類数を決め打ちした場合についてを考えましたが、どの種類を選ぶかがよくわからないため断念しました。そこで、美味しさについてソートすることを考えると、とりあえず $K$ 個選んだパターンは割と高スコアになりそうなので、次にそこからスコアを伸ばせるとしたらどうすればいいか、という部分について考えます。すると、おいしい順に選んでしまっているので、あとは種類数を伸ばすしかありません。よって、最善となる種類の増やし方をあれこれ考えることになり、おいしさが高い寿司順に見ていって、種類をふやせるなら現状の $K$ 個の中から交換、という操作を行うのが最善ということがわかります。あとは、これをバグがないように実装すればいいのですが、無理でした。