原案：kyopro_friends

この問題は、最も美味しさの和が大きな選び方から少しずつ妥協していく過程を、プライオリティーキューを用いてシミュレーションすることで解くことができます。

\(A_i\) をソートすることで、\(A_1\geq A_2 \geq \ldots\) であるとしてよいです。

\(i\) 種類目のクッキーを \(C_i\) 枚選ぶ選び方を、長さ \(N\) の数列により \((C_1,\dots,C_N)\) と表すことにします。

クッキーの選び方 \((C_1,\dots,C_N)\) に対し、\(C_i\neq 0\) であるような \(i\) により \((C_1,\dots,C_{i-1},C_{i}-1,C_{i+1}+1,C_{i+2},\dots)\) の形で書くことのできるものを、一段階妥協した選び方と呼ぶことにします。一段階妥協したクッキーの選び方の美味しさの和は、妥協前のクッキーの選び方の美味しさの和以下なので、美味しさの和の上位から選び方を調べるとき、妥協前の選び方が調べられるまで一段階妥協した選び方が調べられる必要はありません。またどの選び方も \((K,0,0,\ldots)\) から一段階妥協を繰り返すことで到達できます。

よって、美味しさの和と選び方の組をプライオリティーキューで管理しながら、キューから要素を取り出した際に一段階妥協した選び方をキューに追加することで、美味しさの和の順に取り出すことができます。既にキューに追加済みのものを再度追加しないよう注意してください。\(1\) 個の要素をキューから取り出すごとに高々 \(N\) 個の要素がキューに追加されるので \(X\) 番目までを求めることは \(O(NX\log(NX))\) 時間ででき、この問題が解けました。

類題

ABC391F『K-th Largest Triplet』

bonus

\(N\) の上限が \(10^5\) で、他の制約は元と同じ問題を解いてください。