Day 104-Q3 — 行列累乗（Companion Matrix）

2026-07-27 赤色 Master / Phase 8+ ★★★★★★★★☆ 線形漸化式の第N項をO(K^3 log N)で計算

問題

長さ$K$の係数列$c_1,\dots,c_K$と初期値$a_0,\dots,a_{K-1}$が与えられる。数列$a$は$n\ge K$で $a_n = c_1a_{n-1}+c_2a_{n-2}+\cdots+c_Ka_{n-K} \pmod{998244353}$ に従う。整数$N$（最大$10^{18}$）が与えられるので$a_N \bmod 998244353$を求めよ。

直近$K$項をまとめた状態ベクトル $S_i=(a_i,\dots,a_{i-K+1})^\top$ を考えると、$S_i\to S_{i+1}$は固定の $K\times K$ 行列（コンパニオン行列）を掛けるだけの線形変換になる。よって $S_N = M^{N-K+1}\cdot S_{K-1}$ となり、繰り返し二乗法で $O(K^3\log N)$ で計算できる。

入力形式

N K
c_1 c_2 ... c_K
a_0 a_1 ... a_{K-1}

制約

$0 \le N \le 10^{18}$
$1 \le K \le 50$
$0 \le c_i,a_i < 998244353$
$N

入出力例

入力例1

10 3
1 1 1
0 1 1

出力例1

149

$c=(1,1,1)$, $a_0=0,a_1=1,a_2=1$はトリボナッチ数列。$a_3=2,\dots,a_{10}=149$。

概念図: コンパニオン行列と繰り返し二乗法

K=3, c=(1,1,1) のコンパニオン行列 M と M^8 への繰り返し二乗法 M = 111 100 010 1行目=係数c 2,3行目=シフト S0 = a2=1 a1=1 a0=0 → M^8・S0 a10=149 =a9 =a8 指数 8 = 1000(2進) → base を3回二乗するだけで求まる M^1 二乗→ M^2 二乗→ M^4 二乗→ M^8 bit0=0,bit1=0,bit2=0,bit3=1 → resultにM^8だけを掛ける 一般に power のビットが1の桁でのみ result *= base を行う（繰り返し二乗法） 計算量: O(K^3 log N)（K×K行列の乗算 O(K^3) を O(log N) 回）

ヒント（段階的開示）

ヒント1: 方向性
$N\le10^{18}$なので漸化式を$N$回ループする$O(N)$は間に合わない。数の累乗を高速化する「繰り返し二乗法」を、線形変換（行列）に対して適用できないか考えよう。
ヒント2: アプローチ
直近$K$項をまとめたベクトル$S_i=(a_i,\dots,a_{i-K+1})^\top$を考える。$S_i\to S_{i+1}$は固定の$K\times K$行列$M$（1行目=係数、それ以降はシフト）を掛けるだけ。$M$を$N-K+1$乗して初期ベクトルに掛ければ答えが求まる。行列累乗は繰り返し二乗法で$O(K^3\log N)$。
ヒント3: 誘導（コード骨格） 
def mat_pow(mat, power):
    n = len(mat)
    result = 単位行列(n)
    base = mat
    while power > 0:
        if power & 1:
            result = mat_mult(result, base)
        base = mat_mult(base, base)
        power >>= 1
    return result

# M = コンパニオン行列, S0 = [a_{K-1},...,a_0]^T
# S_{N-K+1} = M^{N-K+1} @ S0 の先頭成分が a_N

模範解答 (Python)

import sys

MOD = 998244353


def mat_mult(a, b):
    n = len(a)
    p = len(b)
    m = len(b[0])
    c = [[0] * m for _ in range(n)]
    for i in range(n):
        ai = a[i]
        ci = c[i]
        for kk in range(p):
            if ai[kk] == 0:
                continue
            aik = ai[kk]
            bk = b[kk]
            for j in range(m):
                ci[j] = (ci[j] + aik * bk[j]) % MOD
    return c


def mat_pow(mat, power):
    n = len(mat)
    result = [[1 if i == j else 0 for j in range(n)] for i in range(n)]
    base = mat
    while power > 0:
        if power & 1:
            result = mat_mult(result, base)
        base = mat_mult(base, base)
        power >>= 1
    return result


def solve():
    data = sys.stdin.buffer.read().split()
    idx = 0
    n = int(data[idx]); idx += 1
    k = int(data[idx]); idx += 1
    c = list(map(int, data[idx:idx + k])); idx += k
    a = list(map(int, data[idx:idx + k])); idx += k

    if n < k:
        print(a[n] % MOD)
        return

    m_mat = [[0] * k for _ in range(k)]
    for j in range(k):
        m_mat[0][j] = c[j] % MOD
    for i in range(1, k):
        m_mat[i][i - 1] = 1

    s0 = [[a[k - 1 - i] % MOD] for i in range(k)]

    power = n - k + 1
    m_pow = mat_pow(m_mat, power)
    s_n = mat_mult(m_pow, s0)

    print(s_n[0][0] % MOD)


solve()
計算量: $O(K^3\log N)$。サンプル(トリボナッチ, N=10)で出力149を確認済み（ブルートフォースと一致）。

Step-by-Step 解説

1状態ベクトルの定義
直近$K$項をまとめた$S_i=(a_i,\dots,a_{i-K+1})^\top$を考える。
2コンパニオン行列の構築
1行目に係数$c_1,\dots,c_K$、$i$行目($i\ge2$)は$(i-1)$列目だけ1のシフト行列$M$を作る。
3行列累乗
$S_{N-K+1}=M^{N-K+1}\cdot S_{K-1}$を繰り返し二乗法で$O(K^3\log N)$で計算する。
4境界条件の処理
$N

よくあるミス

ミス原因正しい書き方
$N指数$N-K+1$が負になり破綻する`if n < k: print(a[n]); return`を先頭で処理する
コンパニオン行列の添字をずらす1行目の係数順・シフト行の位置を間違えやすい`m_mat[0][j]=c[j]`、`m_mat[i][i-1]=1`を丁寧に対応させる
行列積の中で`% MOD`を取り忘れる多倍長のまま計算し続け極端に遅くなる各要素の加算のたびに`% MOD`する
結果ベクトルのどの行が$a_N$かを勘違いする$S_i$の成分の並び順（降順）を忘れている`s_n[0][0]`が$a_N$に対応することを確認する

次のステップ

  • 発展: Kitamasa法で行列を持たず$O(K^2\log N)$に高速化する
  • 発展: 係数が未知の場合、先頭$2K$項からBerlekamp-Massey法で係数自体を復元する
  • 次回予告: Manacher's Algorithm — 最長回文部分文字列を$O(N)$で求める

自己評価

自分の回答

気づき・メモ