きたまさ法 (Kitamasa Method)
概要
線形漸化式 \(\displaystyle a_{i+K} = c_0 a_i + c_1 a_{i+1} + ... + c_{i+K-1} a_{i+K-1}\) について \(a_N\) を求める。
解説ブログ: Kitamasa法
計算量
\(O(K^2 \log N)\)
実装
# (0-indexed)
# a[i] = a_i, c[i] = c_i
# a_{i+k} = c_i*a_i + c_{i+1}*a_{i+1} + ... + c_{i+k-1}*a_{i+k-1}
def solve1(m, k, C, A):
MOD = 10**9 + 7
C0 = [0]*k; C1 = [0]*k
if m == 0:
return A[0]
C0[1] = 1
def inc(k, C0, C1):
C1[0] = C0[k-1] * C[0] % MOD
for i in range(k-1):
C1[i+1] = (C0[i] + C0[k-1]*C[i+1]) % MOD
def dbl(k, C0, C1):
D0 = [0]*k; D1 = [0]*k
D0[:] = C0[:]
for j in range(k):
C1[j] = C0[0] * C0[j] % MOD
for i in range(1, k):
inc(k, D0, D1)
for j in range(k):
C1[j] += C0[i] * D1[j] % MOD
D0, D1 = D1, D0
for i in range(k):
C1[i] %= MOD
p = 32
while (m >> p) & 1 == 0:
p -= 1
while p:
p -= 1
dbl(k, C0, C1)
C0, C1 = C1, C0
if (m >> p) & 1:
inc(k, C0, C1)
C0, C1 = C1, C0
return sum(C0[i] * A[i] for i in range(k)) % MOD
# solve1の関数を展開したバージョン
def solve2(m, k, C, A):
MOD = 10**9 + 7
C0 = [0]*k; C1 = [0]*k
if m == 0:
return A[0]
C0[1] = 1
p = 32
while (m >> p) & 1 == 0:
p -= 1
D0 = [0]*k; D1 = [0]*k
while p:
p -= 1
#dbl(k, C0, C1)
D0[:] = C0[:]
for j in range(k):
C1[j] = C0[0] * C0[j] % MOD
for i in range(1, k):
#inc(k, D0, D1)
D1[0] = D0[k-1] * C[0] % MOD
for j in range(k-1):
D1[j+1] = (D0[j] + D0[k-1]*C[j+1]) % MOD
for j in range(k):
C1[j] += C0[i] * D1[j] % MOD
D0, D1 = D1, D0
for i in range(k):
C1[i] %= MOD
C0, C1 = C1, C0
if (m >> p) & 1:
#inc(k, C0, C1)
C1[0] = C0[k-1] * C[0] % MOD
for i in range(k-1):
C1[i+1] = (C0[i] + C0[k-1]*C[i+1]) % MOD
C0, C1 = C1, C0
return sum(C0[i] * A[i] for i in range(k)) % MOD
Verified
-
AtCoder: "Typical DP Contest - T問題: フィボナッチ": source (PyPy3, 1052ms)