問題
数列 $a_1,\dots,a_N$($\mathrm{mod}\ P$)が、次数$d$の線形漸化式 $a_n=c_1a_{n-1}+\dots+c_da_{n-d}$($d$は未知)に従うことが保証されている。与えられた項からこの最小次数の漸化式を復元し、$a_K$($K$は非常に大きくてよい)を$\mathrm{mod}\ P$で求めよ。
入力形式
N K P
a_1 a_2 ... a_N
制約
$2 \le N \le 2000$
$1 \le K \le 10^{18}$
$P$は素数、$10^8 < P < 1.1\times10^9$
入出力例
入力例1
6 10 998244353
1 1 2 3 5 8
出力例1
55
フィボナッチ数列($a_n=a_{n-1}+a_{n-2}$、次数2)の第10項は55。
入力例2
4 5 998244353
2 4 8 16
出力例2
32
等比数列($a_n=2a_{n-1}$、次数1)の第5項は$2^5=32$。
概念図: LFSR(線形帰還シフトレジスタ)としての漸化式
ヒント(段階的開示)
ヒント1: 方向性
漸化式の次数も係数も未知の状態から、与えられた項だけで漸化式を推定する必要がある。次数を1から順に試す方法は非効率。数列を1項ずつ読み、矛盾が起きたときだけ修正するオンラインなアルゴリズムが必要になる。
ヒント2: アプローチ
Berlekamp-Massey法は数列を先頭から読み、現在の漸化式で次項が予測できるか確認し、外れたときだけ最小限の修正を加える。$O(N^2)$で最小次数の漸化式が求まる。漸化式が求まれば、$d\times d$の同伴行列$M$を使い$M^{K-d}$を行列累乗($O(d^3\log K)$)で計算して$a_K$に到達できる。
ヒント3: 誘導(コード骨格)
def berlekamp_massey(S, MOD):
ls, cur = [], []
lf, ld = 0, 0
for i in range(len(S)):
t = sum(cur[j]*S[i-1-j] for j in range(len(cur))) % MOD
if (S[i]-t) % MOD == 0:
continue
if not cur:
cur = [0]*(i+1); lf, ld = i, (S[i]-t)%MOD; continue
k = (S[i]-t) * pow(ld, MOD-2, MOD) % MOD
c = [0]*(i-lf-1) + [k] + [(-k*x)%MOD for x in ls]
# cur との長さ調整・加算、lf/ld/ls の更新は本文参照
cur = c
return cur # a[n] = sum(cur[j]*a[n-1-j])
模範解答 (Python)
import sys
def berlekamp_massey(S, MOD):
ls, cur = [], []
lf, ld = 0, 0
for i in range(len(S)):
t = 0
for j in range(len(cur)):
t = (t + cur[j] * S[i - 1 - j]) % MOD
if (S[i] - t) % MOD == 0:
continue
if not cur:
cur = [0] * (i + 1)
lf, ld = i, (S[i] - t) % MOD
continue
k = (S[i] - t) * pow(ld, MOD - 2, MOD) % MOD
c = [0] * (i - lf - 1) + [k] + [(-k * x) % MOD for x in ls]
if len(c) < len(cur):
c += [0] * (len(cur) - len(c))
for j in range(len(cur)):
c[j] = (c[j] + cur[j]) % MOD
if i - len(cur) >= lf - len(ls):
ls, lf, ld = cur, i, (S[i] - t) % MOD
cur = c
return cur
def mat_mul(A, B, MOD):
n = len(A); m = len(B[0]); k = len(B)
C = [[0] * m for _ in range(n)]
for i in range(n):
Ai, Ci = A[i], C[i]
for l in range(k):
if Ai[l] == 0:
continue
a, Bl = Ai[l], B[l]
for j in range(m):
Ci[j] = (Ci[j] + a * Bl[j]) % MOD
return C
def mat_pow(M, p, MOD):
n = len(M)
result = [[1 if i == j else 0 for j in range(n)] for i in range(n)]
base = M
while p > 0:
if p & 1:
result = mat_mul(result, base, MOD)
base = mat_mul(base, base, MOD)
p >>= 1
return result
def solve():
data = sys.stdin.read().split()
idx = 0
n = int(data[idx]); idx += 1
k = int(data[idx]); idx += 1
P = int(data[idx]); idx += 1
a = [int(data[idx + i]) % P for i in range(n)]
coeffs = berlekamp_massey(a, P)
d = len(coeffs)
if k <= n:
print(a[k - 1])
return
M = [[0] * d for _ in range(d)]
for j in range(d):
M[0][j] = coeffs[j]
for i in range(1, d):
M[i][i - 1] = 1
Mp = mat_pow(M, k - d, P)
v = [a[d - 1 - i] for i in range(d)]
result = 0
for j in range(d):
result = (result + Mp[0][j] * v[j]) % P
print(result % P)
solve()
計算量: Berlekamp-Massey $O(N^2)$、行列累乗 $O(d^3\log K)$。入力例1(フィボナッチ、$d=2$)で$a_{10}=55$、入力例2(等比数列、$d=1$)で$a_5=32$と一致。$K\le N$の境界ケースも配列直読みで正しく処理されることを確認済み。
Step-by-Step 解説
1Berlekamp-Masseyで最小次数の漸化式を推定する
予測が外れたときだけ最小限の修正を加える手法で$O(N^2)$で漸化式を求める。
予測が外れたときだけ最小限の修正を加える手法で$O(N^2)$で漸化式を求める。
2復元した係数から同伴行列を構成する
1行目に係数、それ以外は1つ下にシフトする単位対角線を持つ$d\times d$行列。
1行目に係数、それ以外は1つ下にシフトする単位対角線を持つ$d\times d$行列。
3行列累乗で遠い項へジャンプする
$M^{K-d}$を$O(\log K)$回の行列積で求め、初期状態に1度だけ掛ける。
$M^{K-d}$を$O(\log K)$回の行列積で求め、初期状態に1度だけ掛ける。
4$K\le N$の場合は直接答える
境界条件の処理を忘れない。
境界条件の処理を忘れない。
5全て$\mathrm{mod}\ P$で演算する
除算はフェルマーの小定理による逆元で実現する。
除算はフェルマーの小定理による逆元で実現する。
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| Berlekamp-Masseyの内部変数の役割を混同する | curとlsの意味を自己流に変更する | 標準実装のロジックをそのまま踏襲する |
| 同伴行列の1行目と単位対角線の位置を逆にする | 状態ベクトルの並び順と行列の対応がずれる | 先頭が新しい項の順で統一し1行目=係数、他はM[i][i-1]=1 |
| $K\le N$のケースを行列累乗にかけてしまう | 境界条件チェックの省略 | k<=nなら配列から直接返す分岐を先頭に置く |
| MOD演算の符号処理を誤る | 他言語の癖で剰余が負になると誤解する | Pythonの%は非負を返すためそのままでよい |
次のステップ
- 発展: Kitamasa法(多項式mod演算)で$O(d^2\log K)$、NTTで$O(d\log d\log K)$に高速化する
- 発展: ベクトル値の連立線形漸化式への一般化を考える
- 発展: Library Checkerの「Kth term of Linearly Recurrent Sequence」相当の問題を解いてみる