モノグサプログラミングコンテスト2024（AtCoder Beginner Contest 345）E 問題メモ

モノグサプログラミングコンテスト2024（AtCoder Beginner Contest 345）

E - Colorful Subsequence

• $N$ 個のボールが1列に並び、$i$ 番目のボールは色 $C_i$、価値 $V_i$
• ちょうど $K$ 個のボールを取り除き、同じ色のボールが連続しないようにしたい
• できるか判定し、できる場合は残るボールの価値の最大値を求めよ
• $1 \le N \le 2 \times 10^5$
• $1 \le K \le \min(N,500)$

まず、以下のDPを考える。

• $DP[i,j]=i$ 番目までのボールを考え、$j$ 個取り除き、$i$ 番目のボールは使用したときの、価値の最大値

以下のように求められる。
便宜的に、0番目に色がどれとも異なり価値は0のボールがあり、それは必ず使用すると考える。

  i→
j   | 0  1  2  3  4  5  6
↓--+---------------------
  0 | 0     @
  1 | -        @
  2 | -  -        @
  3 | -  -  -        @ [*]
  4 | -  -  -  -

$DP[6,3]$ を求める場合、@ をつけた $DP[5,3],DP[4,2],...$ の中で、 かつ「$i$ が $C_i \neq C_6$」である中での最大値に、$V_6$ を加えたものとなる。

これを愚直に求めると、$O(NK)$ のマスそれぞれの遷移に $O(K)$ かかり、$O(NK^2)$ となり間に合わない。

更新の際に参照するのは、常にナナメに並ぶ要素であることを利用する。

• $DP2[i,j,k]=DP[i,*]$ までを考え、$(i,j)$ から左斜め上に辿っていった中での、
  • $k=0$: 最大値
  • $k=1$: 最大値の時の最後に置いた色
  • $k=2$: 最大値とは異なる色で、2番目に大きな値

添字 * は、その添字の範囲全体を指すとする。

$DP[6,3]$ を求める場合、$DP2[5,3,*]$ を参照し、 $DP2[5,3,1]$ が $C_6$ と同じなら $DP2[5,3,2]$ を、異なれば $DP1[5,3,0]$ を使えばよい。

そして、新しい $DP[6,*]$ をもって、$DP2[6,*]$ を求める。
この時に $DP[6,3]$ と比較するのは $DP2[5,2,*]$ であり、答えを求める際に使用した $DP2[5,3,*]$ とは1段ずれる点に注意。

これで遷移が $O(1)$ となり、$O(NK)$ で間に合うようになる。

Python3