競プロ備忘録

競プロerの備忘録

ABC345E - Colorful Subsequence

3完しかできなくてひどい目に遭った…
ふと思い返すと、ひどい目に遭ったときしかブログ書いてない気がする。

そんなことはどうでもよくて、Dは解けそうにないし面白い要素もないので、Eをずっと考えていた。
こっちは考えごたえがあってとても面白かったので、コンテスト後だけど自力ACできて嬉しい。

解法

とりあえず愚直なDPを考えてみる。 K \lt 500や実行時間制限5secという不自然な制約があり、ワンチャン通るかもしれないので、書いてみる価値は十分ある。
DPの定義は以下のようになる。

 DP _ {i, j} := 前からボールの使い方を決めることとして、 i番目を使い、そこまでで j個ボールを飛ばしている場合の最高スコア

愚直にDP配列を埋めるのは簡単で、前からボールを見つつ、色が異なる自身の直前 K個以内のボール prevにたいして、 DP _ {i, j} = max( DP _ {i, j}, DP _ {prev, j - (i - prev - 1)} + P _ {i})と更新していけばよい。
この定義だと、 DP _ {i, j}では iが必ず使われることになっているので、最初と最後の場合分けが面倒くさい。なので、最初と最後に番兵として (C, P) = (0, 0)などを挟むとよい。

時間計算量は O(NK ^ {2})となる。一応投げてみると、半分以上のケースでTLEすることがわかる。(提出例:Submission #51323752 - Monoxer Programming Contest 2024(AtCoder Beginner Contest 345)

コードの見た目的に、for prev in (0..i).rev() {....の中身を小さい計算量でできそうなので、ここの高速化を考える。

dp[i][j].max(dp[prev][j-skip] + p[i].1);という箇所に着目すると、p[i].1はずっと同じなので考える必要はない。dp[prev][j-skip]という箇所を考えると、let skip = i - prev - 1;なので、dp[prev][j - (i - prev - 1)]と同じ。この2つの添え字の差を考えると、 prev - j + (i - prev - 1) = i - j - 1であり、prevに依存しなくなることがわかる。
これにすぐ気づくのは難しかったが、DP配列のある箇所の更新の一連の操作を紙でマス目に書いてみると、マス目をナナメに走ることから何となくインデックスの和か差のどっちかが重要なんだろうということが予想できる。

ここまでわかると、DP配列のナナメ方向の累積maxを持てれば良さそうだと予想できる。
ただ、色の制約があるので、実際には累積スコアと最後に使うことを決めたボールの色のペアを、スコアの上位2つ分色違いになるように持てば良い。(以後も、このデータの持ち方やこれを維持する操作を指して「累積max」「max」ということにする)

累積maxということなので、セグメント木が真っ先に思い浮かぶ。しかし、これを書くとTLEする。(提出例:Submission #51353665 - Monoxer Programming Contest 2024(AtCoder Beginner Contest 345)
VecDeqでセグメント木を持つ数を減らしているが、これをやらないとMLEもする。
計算量は O(NKlogK)なので通るだろうと踏んでいたが、よく考えると NKlogKは最大 10 ^ 9近くなるので、まあ通らなくても文句は言えない。

さらにこのコードを睨むと、累積maxを取得する箇所は、常に 0から jまでの累積maxを取得していることがわかる。
中途半端な区間の累積maxが不要であるなら、わざわざセグメント木を使うまでもなく、単に配列で保持することができる。累積maxの更新順序に着目すると、jは小さいほうから更新されるため、値更新を行ったら直前の要素とmaxを取るだけで累積maxの更新もできることがわかる。

ここまでやると、無事ACすることができる。 計算量は O(NK)
提出例は以下の通り。

Submission #51356969 - Monoxer Programming Contest 2024(AtCoder Beginner Contest 345)