連載しょぼんコーダー 1 完結・非完結のDP

連載しょぼんコーダー では, 自分が考えているランダムな競プロテクニックについて解説します.
内容は自分のオリジナルだったりオリジナルじゃなかったりして, ためになるものとためにならないものがあります.

第1回は, 完結・非完結を状態としてとらえることで, 遷移が前の項のみに依存するDP(動的計画法)を考えます.

完結された状態が $1$ 個, 未完結の状態が $n$ 個あるとして, $n+1$ 個の状態を持つDPを考えます.
$DP(i, j)$ の $i$ には, $i$ 番目まで見たという情報を入れ, $j ~ (0 \leq j \leq n)$ には状態を入れます.
こうすると, 完結状態から再び未完結状態に遷移させることを, 未完結状態から完結状態に遷移させるのと同じように実装できます.

例題 $1$ （想定Diff 1000）
$N$ 項の数列 $A = (A_1, A_2, \cdots, A_N)$ が与えられます.
数列 $A$ から連続するちょうど $2$ 個の項を重複せずに選ぶことを繰り返します.
その総和が $M$ 以下になる選び方は何通りありますか. （選ぶ順番は考慮しません.）
（998244353で割った余りを答えてください.）
制限時間：2秒, メモリ制限：1GB
制約： $2 \leq N \leq 3000, ~ 1 \leq M \leq 3000, ~ 0 \leq A_i \leq 10^9$
たとえば, $N=7, M=11, A=(3, 1, 4, 1, 5, 9, 2)$ のとき, 選び方は
$(1)~ (3, 1), 4, 1, 5, 9, 2$
$(2)~ 3, (1, 4), 1, 5, 9, 2$
$(3)~ 3, 1, (4, 1), 5, 9, 2$
$(4)~ 3, 1, 4, (1, 5), 9, 2$
$(5)~ 3, 1, 4, 1, 5, (9, 2)$
$(6)~ (3, 1), (4, 1), 5, 9, 2$
$(7)~ (3, 1), 4, (1, 5), 9, 2$
$(8)~ 3, (1, 4), (1, 5), 9, 2$
の8通りです.

（注）次のように入力されます. 答えは一行で出力してください.

N M
A_1 A_2 ... A_N

天才解法として, $N-1$ 項の数列 $B = (A_1 + A_2, A_2 + A_3, \cdots, A_{N-1} + A_N)$ について, $B$ における部分列であって, どの項も隣り合わないもののうち, その値の総和が $M$ 以下になるようなものの選び方に帰着できます.
$DP(i, j)$ を, $i$ ( $0\leq i \leq N-1$ )番目まで見た総和が $j$ ( $0\leq i \leq M$ )である場合の数とすると,

初期値は $DP(0, 0)=1, 他=0$ で,
選ぶ場合は $i-2$ 番目のときの状態から選ぶ必要があるので, $DP(i-2, j-B_i)$ から遷移します.（ $j-B_i \lt 0$ ならば遷移しません.）
選ばない場合は $i-1$ 番目のときの状態からそのまま移るので, $DP(i-1, j)$ から遷移します.

結局, $DP(i, j) = DP(i-2, j-B_i) + DP(i-1, j)$ とできるので, $DP(N-1, 0) + DP(N-1, 1) + \cdots + DP(N-1, M)$ を出力すればokです.

mod = 998244353

N, M = map(int,input().split())
A = list(map(int,input().split()))
B = [A[i] + A[i+1] for i in range(N-1)]

DP = [[0] * (M+1) for j in range(N)]
DP[0][0] = 1
for i in range(N-1):
    for j in range(M+1):
        DP[i+1][j] += DP[i][j]
        if j - B[i] >= 0:
            DP[i+1][j] += DP[max(0, i-1)][j - B[i]]
        DP[i+1][j] %= mod

#print(*DP, sep="\n")

ans = mod-1
for i in range(M+1):
    ans += DP[N-1][i]
    ans %= mod

print(ans)

いまのは天才でした. 次に, 完結か未完結か保持するDPでこの問題を解いてみます. $B$ を構成する必要はないけれども, 少し計算量が増えます.
$DP(i, j, k)$ ： $i$ ( $0 \leq i \leq N$ )番目まで見て, 総和が $j$ ( $0\leq i \leq M$ )であり, $k$ （ $k \in {0, 1}$ ）0ならば完結, 1ならば未完結 であるものの場合の数とします.
完結とは,選び方として正しい状態を意味して, 未完結とは, 連続する2枚のうち前者を選んでいて, 後者はまだ選んでいない状態とします.（前者を選ぶと, 後者も選ぶ必要があるので）

初期値は $DP(0, 0, 0)=1, 他=0$ で,
新しく選ぶ場合は重複しないように完結している状態から選ぶので, $DP(i-1, j-A_i, 0)$ から遷移します.（ $j-A_i \lt 0$ ならば遷移しません.）
前回選んだ場合, 状態は未完結で $k=1$ に残っているので, 必ず次も選んで $k=0$ に戻します. よって $DP(i-1, j-A_i, 1)$ から遷移します.（ $j-A_i \lt 0$ ならば遷移しません.）
前回選んでおらず, 今回も選ばない場合, $i-1$ 番目のときの状態からそのまま移るので, $DP(i-1, j, 0)$ から遷移します.

$k=1$ の方： $DP(i, j, 1) = DP(i-1, j-A_i, 0)$
$k=0$ の方： $DP(i, j, 0) = DP(i-1, j-A_i, 1) + DP(i-1, j, 0)$
として, 結局 $DP(N, 0, 0) + DP(N, 1, 0) + \cdots + DP(N, M, 0)$ を出力すればokです.

mod = 998244353

N, M = map(int,input().split())
A = list(map(int,input().split()))

DP = [[[0] * 2 for j in range(M+1)] for i in range(N+1)]
DP[0][0][0] = 1
for i in range(N):
    for j in range(M+1):
        DP[i+1][j][0] += DP[i][j][0]
        if j - A[i] >= 0:
            DP[i+1][j][1] += DP[i][j - A[i]][0]
            DP[i+1][j][0] += DP[i][j - A[i]][1]
        DP[i+1][j][0] %= mod
        DP[i+1][j][1] %= mod

#print(*DP, sep="\n")

ans = mod-1
for i in range(M+1):
    ans += DP[N][i][0]
    ans %= mod

print(ans)

遷移が前の項にしか影響しないことを利用すると, 次のように定数倍高速化できます.

mod = 998244353

N, M = map(int,input().split())
A = list(map(int,input().split()))

DP = [[0] * 2 for j in range(M+1)]
NDP = [[0] * 2 for j in range(M+1)]
DP[0][0] = 1
for i in range(N):
    for j in range(M+1):
        NDP[j][0] = 0
        NDP[j][1] = 0
    for j in range(M+1):
        NDP[j][0] += DP[j][0]
        if j - A[i] >= 0:
            NDP[j][1] += DP[j - A[i]][0]
            NDP[j][0] += DP[j - A[i]][1]
        NDP[j][0] %= mod
        NDP[j][1] %= mod
    for j in range(M+1):
        DP[j][0] = NDP[j][0]
        DP[j][1] = NDP[j][1]

#print(*DP, sep="\n")

ans = mod-1
for i in range(M+1):
    ans += DP[i][0]
    ans %= mod

print(ans)