トヨタ自動車プログラミングコンテスト2024#1（AtCoder ABC337） F問題 Usual Color Ball Problems

めちゃくちゃ苦労したけど、コンテスト外でやっと解けたので記録してくおく。

最初に考えたこと。
差分計算したい。
更新時に未使用の色のうち、インデックスが一番小さいものを採用すればよさそう。
ある色が採用されたとき、いくつのボールが利用されるかは事前計算しておく。
色をインデックスとしたセグ木に、各色の一番小さいインデックスをいれて更新していく。使用中の色にはInfを入れておく。
この方針はだいたいのケースでうまくいくんだけど、考慮漏れがあって、一つの色が複数の箱に入る場合にWAになる。
そこでさらに発展させて、箱が無限にあったとき、各色が何個の箱に入ることができるかと、n箱目では何個のボールが入っているかを事前に計算しておくことで、いい感じにセグ木の更新と求める合計を計算することができる。
と書くのは簡単だけど、実装は結構大変だった。コンテスト時間ないに解くのは今のところ絶望的。

擬似的な説明を書く。
各色のボールがn箱目で利用されたときの利用されるボールの数（最大K)：ball[color][n]
各色のボールのindexを管理する配列（予めCを2周分にしておく）。各色の左からt番目。存在しないときはInfを返す。 : index[color][t]
各色について未使用のボールのうち最も左のindexを管理するセグ木。存在しないときはinfをいれる。: segtree[color]
各色について何箱利用しているか : count[color]
初期状態の答え: sum

ballはpop()で先頭を消せるようにリバースしておく。

for _ in 0..M{ //M箱分処理する
    //いい感じに初期状態とsumを作る
}

let mut ans=vec![sum];
for i in 0..N-1{
    let preColor = C[i];
    count[preColor]-=1;    
    index[preColor].pop();
    sum-=ball[preColor][count[preColor]];
    
    segtree.update(preColor, index[preColor].last());
    
    let j = segtree.all_prod();
    let nextColor = C[j];
    sum+=ball[nextColor][count[nextColor]];
    count[nextColor]+=1;

    let L= index[nextColor].len();
    let n = index[nextColor][L - count[nextColor]*K -1];
    //未使用の次のインデックスは現在（末尾）から使用している箱*K番目    
    segtree.update(nextColor, n);

    ans.push(sum);    

}

以上を実装してAC。