ボールを入れ終わる前にも箱には値が表示されていることにします。 はじめ、\(K\) 個の箱に \(0\) が表示されています。

\(0\) が書かれたボールの個数を \(X\) と表すことにします。 最大で \(\min(K,X)\) 個の箱は \(1\) 以上の値を表示させることができます。一方で、少なくとも \(K-\min(K,X)\) 個の箱は \(0\) 以外の値を表示させることはできません。

\(0\) が書かれたボールを箱に入れた後、 \(1\) が表示されていない箱を捨て、箱に表示された値と残ったボールに書かれた値を \(1\) 減らすと元の問題に戻ることがわかります。 また、最終的な答えに残った箱の個数の値だけ寄与することもわかります。ここから、箱はできる限り多く残したほうがよいことが示せます。

結局、書かれた値が小さいボールから順に箱に入れていき、可能な限り箱に表示される値を変化させればよいです。 これは \(O(N \log N)\) や \(O(N)\) で実行可能で \(N \leq 3 \times 10^5\) より十分高速です。