ABC 401 C - K-bonacci

↑問題の詳細は上記のリンクを参照ください！
※興味持った方は毎週土曜日21:00からのABCコンテストにratedで参加してみてください⭐️

問題文

正整数 N, K が与えられます。長さ N+1 の数列 A = (A₀, A₁, ..., Aₙ) の各要素の値を、以下の方法で定義します。

• 0 ≦ i < K のとき、Aᵢ = 1
• K ≦ i のとき、Aᵢ = A_{i-K} + A_{i-K+1} + ... + A_{i-1}

A_N を 10の9乗（10⁹）で割ったあまりを求めてください。

制約

• 1 ≦ N, K ≦ 10⁶
• 入力される数値はすべて整数

解法ポイント

こんな感じに整理するとK個分の数列Aを配列で管理してあげると解けそうです。

NG例1

from collections import deque

MOD = 10**9
N,K = map(int, input().split())

nums = [1 for _ in range(K)]
dq = deque()
dq.extend(nums)

for _ in range(N-K+1):
  dq.append(sum(dq) % MOD) # O(N*N)になっちゃう。❌制約がN<=10**6なのでMAX 10**12
  dq.popleft()

print(dq[-1])

体感python(pypy)だと10**10超えるとタイムアウトでTLE(誤答)になる感覚です

NG例2

NG例1からO(N*N)からO(N)にどうやってできるか試行錯誤すると
実は、以下のことに気づくとO(N)で計算可能なことに気づきます。

i(>=K+1)の場合は、A[i] = A[i-1]*2 - A[i-K-1]

from collections import deque

MOD = 10**9
N,K = map(int, input().split())

nums = [1 for _ in range(K)]
dq = deque()
dq.extend(nums)

total = sum(dq)

for _ in range(N-K+1):
  dq.append(total % MOD)
  total = (total *2 - dq.popleft()) #ここにMODを置かないと桁がデカくなりすぎてTimeOutする

print(dq[-1])

ACサンプル(正解)

ということで、totalの再計算にもMODを入れるとAC(合格)します。

from collections import deque

MOD = 10**9
N,K = map(int, input().split())

nums = [1 for _ in range(K)]
dq = deque()
dq.extend(nums)

total = sum(dq)

for _ in range(N-K+1):
  dq.append(total % MOD)
  total = (total *2 - dq.popleft()) % MOD

print(dq[-1])