概要
paiza x Qiita コラボキャンペーンの問題「十億連勝(paizaランク S 相当)」を $O(N)$ で解きます。
O(N^2) 解法
公式解説にあるように、DP テーブルのキーとして
(どこまで見たか, 現在の連勝数, ちょうど X 連勝達成フラグ)
を持ちます。
v = dp[i-1][num][flag]からの遷移は次のようになります。
- $A_i$ 連勝すると$X$連勝をオーバーするとき（$\mathrm{num} + A_i > X$）は
-
dp[i][0][flag] += v*(X-num)（$X$ 連勝未満で負ける） -
dp[i][0][true] += v（$X$ 連勝ちょうどで負ける）
-
- 連勝がオーバーしないときは
-
dp[i][num+A[i]][flag] += v（全勝する） -
dp[i][0][flag] += v*A[i]（どこかで負ける）
-
DP テーブルを辞書で管理すれば、管理するDPテーブルの個数は $O(N^2)$ 個です（計算量は辞書のオーダーがかかります）。
また連勝数の代わりに「最後に負けたステージ番号」を管理すれば辞書を使う必要がなくなり、計算量自体が $O(N^2)$になります。
O(N) 解法
いわゆる in-place DP （インラインDP）です。
ここではDP テーブルを deque で管理し、deque の中で連勝数が単調増加になるようにします。
（「最後に負けたステージ番号」を管理というアイデアでも問題なく実装できます）
これにより次のメリットがあります。
- 負けて $0$ 連勝になるとき、
appendleftにより単調性を保って値が追加できる - 連勝がオーバーするかどうかを deque の末尾のみを見て判定できる
連勝がオーバーするものがある限り、deque から pop して DP テーブルを更新します（実際には $0$ 連勝の場合の数をカウントしていって、最後に更新します）。
さらに連勝がオーバーしないときの DP の遷移をまとめて行います。次の値を管理します。
-
ng_total,ok_total: dp テーブルの中の 達成していない/達成した の和 -
offset: 見かけの連勝数と実際の連勝数の差分
すると、連勝がオーバーしないときの DP の遷移は以下のように表され、遷移をまとめて $O(1)$ で実行できます。
- どこかで負ける場合:
ng_total,ok_totalの値に倍率 $A_i$ をかけて $0$ 連勝の DP の値を更新する - 全勝する場合:
offsetをずらす（すべての連勝数に $A_i$ を足すことに相当する）
Deque への追加は $N+1$ 回、削除も高々 $N+1$ 回なので、計算量はならし $O(N)$ です。
from collections import deque
MOD = 10**9
N,X = map(int,input().split())
"""
dp テーブル: (key=見かけの連勝数, ng=達成していない場合の数, ok=達成した場合の数)
不変条件: ng_total, ok_total: dp テーブルの中の ng/ok の和
不変条件: key + offset = 実際の連勝数
"""
dp = deque()
dp.append((0,1,0))
offset = 0
ng_total, ok_total = 1,0
for _ in range(N):
A = int(input())
ng_new = ok_new = 0 # 途中で負けて 0 連勝になる場合の数を保持する変数
# 最大連勝で X+1 連勝以上する場合
while dp and dp[-1][0] + offset + A > X:
key, ng_old, ok_old = dp.pop()
# 不変条件の維持
ng_total = (ng_total - ng_old) % MOD
ok_total = (ok_total - ok_old) % MOD
# X-1 連勝以下で負ける
num = X - (key + offset)
ng_new = (ng_new + num * ng_old) % MOD
ok_new = (ok_new + num * ok_old) % MOD
# X 連勝で負ける
ok_new = (ok_new + ok_old + ng_old) % MOD
# 最大連勝でも X 連勝以下の場合
#途中で負ける
ng_new = (ng_new + A * ng_total) % MOD
ok_new = (ok_new + A * ok_total) % MOD
#全勝する
offset += A
# 0 連勝の追加と不変条件の維持
dp.appendleft((-offset, ng_new, ok_new))
ng_total = (ng_total + ng_new) % MOD
ok_total = (ok_total + ok_new) % MOD
ans = 0
for key, ng, ok in dp:
ans += (ng if key + offset == X else 0) + ok
print(ans%MOD)