Toyota Programming Contest 2024#1（AtCoder Beginner Contest 337）F問題メモ

Toyota Programming Contest 2024#1（AtCoder Beginner Contest 337）

GよりむずいF問題

F - Usual Color Ball Problems

• $N$ 個のボールが1列に並び、それぞれの色は $C_i$
• $M$ 個の、1つにつきボールを $K$ 個まで入れられる空の箱がある
• $i=1,2,...,N$ について、以下を求めよ
  • ボールを、$i$ 番目が先頭になるように並べなおす。$1,...,i-1$ 番目のボールはそのままの順番で末尾にもっていく
  • 先頭から順に、以下の要領でボールを箱に入れる（または捨てる）
    • ボールの色を $c$ として、既に色 $c$ が入っていて $K$ 個未満の箱があれば、そこに入れる
    • なければ、空の箱があれば適当に1つ選んで入れる
    • それもなければ捨てる
  • 最終的に、箱に入ったボールの個数を求めよ
• $1 \le M,K \le N \le 2 \times 10^5$

$i=1$ の時を求めて、1つずらすだけだとほとんどは同じ状態になるだろうので、$i=2,3,...$ の答えは差分で計算したい。

しかし、実際に箱に入るボールは列中で飛び飛びに現れるので、 区間を1つずらしたときに、どの色の箱が無くなり、どのボールが新たな箱を作るのか、結構変化してしまい、効率的に管理しづらい。。

M=3 K=2

Ci         1 2 3 5 2 5 4
箱に入る?  o o o x o x x

Ci           2 3 5 2 5 4 1
箱に入る?    o o o o o x x

Ci             3 5 2 5 4 1 2
箱に入る?      o o o o x x o

Ci               5 2 5 4 1 2 3
箱に入る?        o o o o x o x

重要な考察として、最終的に色 $c$ のボールが箱に入る個数というのは、以下の小さい方になる。
（これに気付かないと進めないのに、なかなか気付きづらい……）

• 色 $c$ のボールが入った箱数 $\times K$
• 色 $c$ のボールの個数 $g(c)$

なので、先頭からボールを入れていくようにシミュレーションする際、 「全ての箱の色が確定したタイミング」、つまり「最後の箱に、1個ボールが入ったタイミング」で、 全体の箱に入ったボールの個数が確定するので、打ち切ってよい。

そこまでのボールは、必ず何かしらの箱に入っていることが保証される。
飛び飛びでは無くなるので、$i$ をずらすたび、途中から探索を行える。

Ci         1 2 3 5 2 5 4
箱に入る?  o o o - - - -      ^ : 最後の箱にボールが入ったタイミング
               ^              - : 未調査
Ci           2 3 5 2 5 4 1
箱に入る?    o o o - - - -
                 ^
Ci             3 5 2 5 4 1 2
箱に入る?      o o o - - - -
                   ^
Ci               5 2 5 4 1 2 3
箱に入る?        o o o o - - -
                       ^

尺取法の要領で、区間の左右端のポインタを管理しておき、区間を1つずらした時、

• 左端のボールを取り除く。取り除いた前後で箱が1つ減れば、差分更新する
• それによって箱が1つ空になった場合、右端からボールを追加していく。次に新たに箱が作られたら差分更新し、終わる

高々2色について、答えを更新すればよい。

Python3