モノグサプログラミングコンテスト2024（AtCoder Beginner Contest 345）E 問題メモ
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を考える。
以下のように求められる。
便宜的に、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)$ から左斜め上に辿っていった中での、
添字 * は、その添字の範囲全体を指すとする。
$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
import os
import sys
import numpy as np
def solve(inp):
n, k = inp[:2]
ccc = inp[2::2]
vvv = inp[3::2]
INF = 1 << 60
dp1v = np.full(k + 1, -INF, np.int64)
dp1c = np.full(k + 1, -INF, np.int64)
dp2v = np.full(k + 1, -INF, np.int64)
dp1v[0] = 0
for i in range(n):
c = ccc[i]
v = vvv[i]
dp_tmp = np.full(k + 1, -INF, np.int64)
for j in range(min(i + 1, k + 1)):
dp_tmp[j] = dp1v[j] if dp1c[j] != c else dp2v[j]
dp_tmp[j] += v
if i < k:
dp_tmp[i + 1] = 0
for j in range(min(i + 1, k), 0, -1):
# print(' ', j, dp1v[j - 1], dp2v[j - 1], dp_tmp[j])
if dp1v[j - 1] < dp_tmp[j]:
if dp1c[j - 1] == c:
dp2v[j] = dp2v[j - 1]
dp1v[j] = dp_tmp[j]
dp1c[j] = c
else:
dp2v[j] = dp1v[j - 1]
dp1v[j] = dp_tmp[j]
dp1c[j] = c
elif dp1c[j - 1] != c and dp2v[j - 1] < dp_tmp[j]:
dp2v[j] = dp_tmp[j]
dp1v[j] = dp1v[j - 1]
dp1c[j] = dp1c[j - 1]
else:
dp2v[j] = dp2v[j - 1]
dp1v[j] = dp1v[j - 1]
dp1c[j] = dp1c[j - 1]
dp1v[0] = dp_tmp[0]
dp1c[0] = c
if dp1v[k] < 0:
return -1
else:
return dp1v[k]
SIGNATURE = '(i8[:],)'
if sys.argv[-1] == 'ONLINE_JUDGE':
from numba.pycc import CC
cc = CC('my_module')
cc.export('solve', SIGNATURE)(solve)
cc.compile()
exit()
if os.name == 'posix':
# noinspection PyUnresolvedReferences
from my_module import solve
else:
from numba import njit
solve = njit(SIGNATURE, cache=True)(solve)
print('compiled', file=sys.stderr)
inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)
