非想定解あるある、制約のせいで逆に迷った。
問題概要
考えたこと
今回は 1-indexed で考えた方が圧倒的にわかりやすいのでそのように統一する。
箱の色容量
色 $ i $ のボールを $ j $ 番目の箱に $ j $ 個入れる操作を「 $ 1 $ 色入れる」と定義する。
このとき、色 $ i $ のボールを箱 $ j $ に $ ( i \cdot j ) $ 個入れる操作は「 $ i $ 色入れる」ことに等しい。
箱 $ j $ の容量は $ B _ j $ なので、箱 $ j $ に入れることができる色の総数（＝色容量）を $ C _ j $ とおくと、
$$ C _ j = \frac { B _ j } { j } \notag $$
である。
ボールを入れる
色容量 $ C _ j $ をなるべく余らせないようにしたいので、$ C _ j $ の大きい箱から優先的にボールを入れていきたい。
この時点で箱 $ j $ の容量 $ B _ j $ を考慮するのは面倒なので、一旦 $ B _ j $ を無視してボールを割り当てて、後で均すこととする。
色 $ i $ のボールを入れる箱の番号の総数の上限は $ \dfrac { A _ i } { i } $ であるから（"箱の色容量"と同様の議論）、これを超えないように $ C _ j $ の大きい順に箱を選ぶ。このとき、各色 $ i $ に対して、
- ボールがちょうど $ ( i \cdot j ) $ 個入った箱
- ボールが $ ( i \cdot j ) $ 個未満入った箱（高々 $ 1 $ 箱）
- ボールが入らなかった箱
が存在する。1. については imos 法や双対セグメント木などで色の総数を管理すれば、箱 $ j $ に割り当てたボールの総数（ $ D _ j $ とする）を復元できる。
ボールを均す
$ C _ j $ の大きい箱から優先的にボールを入れた都合上、$ C _ l < C _ r $ を満たす箱 $ l , r $ について箱 $ l $ から箱 $ r $ へボールを移すことはできない。箱 $ l $ に入っているボールはどの色についても箱 $ r $ に満量入っているためである。
したがって、箱 $ r $ から箱 $ l $ へボールを移すことのみを考えればよい。
結論から言うと、箱 $ r $ から箱 $ l $ へボールを移す際に色の制約は実質的に無視することができて、箱 $ j $ の容量 $ B _ j $ のみに注意して可能な限りボールを移してしまえばよい。
上図のようにボールを移すと、箱 $ l $ に移動した部分は高さ（＝ $ 1 $ 色あたりのボールの個数）が圧縮されるので、色 $ i $ のボールは $ ( i \cdot l ) $ 個以下となっている。
ボールと箱容量の役割交換
計算量は、
- 箱を $ C _ j $ の降順にソートする部分：$ O ( M \log M ) $
- ボールを入れる部分：二分探索により $ O ( N \log M ) $
- ボールを均す部分： $ O ( M ) $
であり、全体で $ O ( ( N + M ) \log M ) $ である。
ボールと箱容量のマッチング問題であるため、ボールを箱容量に、箱容量をボールに読み替えてもよい。
そうした場合、$ O ( ( N + M ) \log N ) $ に収まる。
