Matrix Exponentiation, in depth

The base technique in brief

A linear recurrence of order k defines each term from the previous k: a(n) = c1 a(n-1) + c2 a(n-2) + ... + ck a(n-k). Collect the last k terms into a state vector v(n) = [a(n), a(n-1), ..., a(n-k+1)]. One step of the recurrence is then a fixed k-by-k matrix M with the coefficients in the first row and a shifted identity below it, so v(n+1) = M v(n), and v(n) = M^(n-k+1) v(k-1).

Raising M to a power uses the same square-and-multiply idea as modular exponentiation: M^13 = M^8 M^4 M^1, and each M^(2^i) is the square of the previous one. Each multiplication costs k^3 scalar operations, so the whole computation costs O(k^3 log n). Everything below either changes what goes into the state vector, or replaces the k^3 with something cheaper.

MOD = 1_000_000_007

def mat_mul(A, B, mod=MOD):
    n, m, p = len(A), len(B), len(B[0])
    C = [[0] * p for _ in range(n)]
    for i in range(n):
        Ai, Ci = A[i], C[i]
        for t in range(m):
            a = Ai[t]
            if a:
                Bt = B[t]
                for j in range(p):
                    Ci[j] = (Ci[j] + a * Bt[j]) % mod
    return C

def mat_pow(M, e, mod=MOD):
    k = len(M)
    R = [[int(i == j) for j in range(k)] for i in range(k)]   # identity
    while e:
        if e & 1:
            R = mat_mul(R, M, mod)
        M = mat_mul(M, M, mod)
        e >>= 1
    return R

Augmenting the state vector

The trick that handles most non-standard recurrences is to add components to the state vector that evolve linearly on their own. Anything that can be computed from the previous state by a linear formula can ride along in the matrix.

• A constant term. For a(n) = 2a(n-1) + 5, add a component that is always 1. Its row in the matrix keeps it at 1, and the constant 5 sits in the first row in that component's column.
• A polynomial in n. For a term like n, add components n and 1, with n+1 = n + 1 as the update. For n squared, add n squared, n and 1, using (n+1)^2 = n^2 + 2n + 1. A polynomial of degree d needs d+1 extra components, with binomial coefficients in their rows.
• An exponential term. For a term like 5 times 3^n, add a component holding 3^n whose row multiplies it by 3 each step.

The rule is: every component must be updated from the previous vector using fixed coefficients. Anything that needs a product of two state values, such as a(n-1) times a(n-2), or a coefficient that depends on n, breaks linearity and cannot be expressed this way. Coefficients that depend on n are the most common trap; a term like n times a(n-1) is not linear in the state.

Prefix sums and geometric series

Many problems ask for S(n) = a(0) + a(1) + ... + a(n), not a(n) itself. Add S to the state. Since S(n+1) = S(n) + a(n+1), and a(n+1) is already a linear combination of the state, the row for S is the row for a plus a 1 in the column of S. The matrix grows by one, and the answer appears in the vector at no extra asymptotic cost.

A related problem is summing powers of a whole matrix: I + M + M^2 + ... + M^(n-1). This arises when you need the number of walks of every length up to n in a graph, not just length exactly n. Build the block matrix A = [[M, I], [0, I]], twice the size. Its powers are A^n = [[M^n, I + M + ... + M^(n-1)], [0, I]], which you can check by multiplying A by itself once. So one exponentiation of a 2k-by-2k matrix gives both the power and the geometric sum, without needing a modular inverse of M minus I, which may not exist.

Worked example: a recurrence with extra terms

Take a(n) = 2a(n-1) + 3a(n-2) + n + 1 with a(0) = 1 and a(1) = 1, and suppose we want both a(n) and S(n) = a(0) + ... + a(n). The state is v(n) = [a(n), a(n-1), n+1, 1, S(n)]. To step to v(n+1): a(n+1) = 2a(n) + 3a(n-1) + (n+1) + 1, so its row is [2, 3, 1, 1, 0]. The second component copies a(n), giving [1, 0, 0, 0, 0]. The counter becomes (n+1) + 1, giving [0, 0, 1, 1, 0]. The constant stays 1. The sum adds a(n+1), giving [2, 3, 1, 1, 1].

Check it by hand. Start from v(1) = [1, 1, 2, 1, 2]. One step: a(2) = 2 + 3 + 2 + 1 = 8 and S(2) = 2 + 8 = 10, so v(2) = [8, 1, 3, 1, 10]. Next: a(3) = 16 + 3 + 3 + 1 = 23, S(3) = 33. Next: a(4) = 46 + 24 + 4 + 1 = 75, S(4) = 108. Direct evaluation of the recurrence gives the same numbers, which is the test you should always run before trusting a matrix for n = 10^18.

def solve(n, mod=MOD):
    if n == 0:
        return 1, 1                       # a(0), S(0)
    M = [[2, 3, 1, 1, 0],
         [1, 0, 0, 0, 0],
         [0, 0, 1, 1, 0],
         [0, 0, 0, 1, 0],
         [2, 3, 1, 1, 1]]
    v1 = [1, 1, 2, 1, 2]                  # v(1)
    P = mat_pow(M, n - 1, mod)
    v = [sum(P[i][j] * v1[j] for j in range(5)) % mod for i in range(5)]
    return v[0], v[4]

def brute(n):
    a = [1, 1]
    for i in range(2, n + 1):
        a.append(2 * a[-1] + 3 * a[-2] + i + 1)
    return a[n], sum(a[: n + 1])

assert all(solve(n) == tuple(x % MOD for x in brute(n)) for n in range(30))

Kitamasa: from k^3 to k^2

For a plain order-k recurrence, the matrix approach does k^3 work per step, but the matrix has a lot of structure that it ignores. Kitamasa's method uses it. By the Cayley-Hamilton theorem, M satisfies its own characteristic polynomial P(x) = x^k - c1 x^(k-1) - ... - ck. So M^n equals R(M), where R(x) is x^n reduced modulo P(x), a polynomial of degree below k. Applying that to the initial terms gives a(n) = r0 a(0) + r1 a(1) + ... + r(k-1) a(k-1), where the r are the coefficients of R.

Computing x^n mod P(x) by square-and-multiply needs only polynomial multiplication and reduction, each O(k^2) with schoolbook arithmetic, so the total is O(k^2 log n) instead of O(k^3 log n). For k = 500 that is about 500 times faster. With FFT-based polynomial multiplication it drops further to O(k log k log n), though that only pays off for large k.

def kitamasa(c, init, n, mod=MOD):
    # a(n) = c[0] a(n-1) + ... + c[k-1] a(n-k); init = a(0..k-1)
    k = len(c)
    if n < k:
        return init[n] % mod

    def mul(p, q):                        # (p * q) mod P(x)
        r = [0] * (2 * k - 1)
        for i, pi in enumerate(p):
            if pi:
                for j, qj in enumerate(q):
                    r[i + j] = (r[i + j] + pi * qj) % mod
        for d in range(2 * k - 2, k - 1, -1):   # x^d = sum c_i x^(d-i)
            t = r[d]
            if t:
                for i in range(1, k + 1):
                    r[d - i] = (r[d - i] + t * c[i - 1]) % mod
        return r[:k]

    result = [1] + [0] * (k - 1)                 # x^0
    base = [0, 1] + [0] * (k - 2) if k > 1 else [c[0] % mod]   # x mod P
    while n:
        if n & 1:
            result = mul(result, base)
        base = mul(base, base)
        n >>= 1
    return sum(r * a for r, a in zip(result, init)) % mod

Kitamasa needs a pure homogeneous recurrence. If you have constant or polynomial terms, either fold them in by augmentation first, which raises k, or convert the recurrence: for example, a(n) = 2a(n-1) + 5 implies a(n) - a(n-1) = 2a(n-1) - 2a(n-2), giving the homogeneous a(n) = 3a(n-1) - 2a(n-2).

Finding the recurrence with Berlekamp-Massey

Sometimes you can compute the first few hundred terms of a sequence by brute force or dynamic programming but have no closed recurrence. If the sequence does satisfy a linear recurrence of order k over a prime field, the Berlekamp-Massey algorithm finds the shortest one from 2k terms in O(k^2) time. It suits counting problems, such as tilings of a fixed-width strip, where a recurrence must exist but is tedious to derive.

def berlekamp_massey(s, mod=MOD):
    # returns c with s[n] = sum c[i] * s[n-1-i] for the shortest such recurrence
    C, B = [1], [1]
    L, m, b = 0, 1, 1
    for n in range(len(s)):
        d = s[n]
        for i in range(1, L + 1):
            d = (d + C[i] * s[n - i]) % mod
        if d == 0:
            m += 1
            continue
        coef = d * pow(b, mod - 2, mod) % mod
        T = C[:]
        C += [0] * max(0, len(B) + m - len(C))
        for i in range(len(B)):
            C[i + m] = (C[i + m] - coef * B[i]) % mod
        if 2 * L <= n:
            L, B, b, m = n + 1 - L, T, d, 1
        else:
            m += 1
    return [(-x) % mod for x in C[1 : L + 1]]

Feed it more terms than you think you need, then verify: the recurrence it returns must reproduce terms you held back. If the order it finds keeps growing as you add terms, the sequence is probably not linearly recurrent, or not with an order you can afford. Combined with Kitamasa, this gives a powerful pipeline: brute-force 2k terms, recover the recurrence, jump to term 10^18. Classic dynamic programs like coin change counting with fixed coin sets produce exactly this kind of sequence.

Many queries on one matrix

If you must answer q queries for different n with the same matrix, recomputing M^n each time costs O(q k^3 log n). Precompute instead M, M^2, M^4, ..., M^(2^60) once, costing O(k^3 log n). Then answer each query by applying the needed powers to the vector, one at a time: a matrix times a vector costs only k^2. Each query is O(k^2 log n), and the matrices are never multiplied together again.

POW = [M]
for _ in range(60):
    POW.append(mat_mul(POW[-1], POW[-1]))

def apply(n, v):
    i = 0
    while n:
        if n & 1:
            v = [sum(POW[i][r][j] * v[j] for j in range(len(v))) % MOD
                 for r in range(len(v))]
        n >>= 1
        i += 1
    return v

Failure modes

• Off-by-one exponents. Starting from v(1) means the exponent is n - 1, not n. Always compare against brute force for small n.
• Row versus column convention. Mixing row-vector and column-vector conventions transposes the matrix silently. Pick one and write it in a comment.
• Overflow. In languages with fixed-width integers, reduce after every multiply-add, and use 128-bit intermediates or Montgomery reduction when the modulus is near 2^63.
• Non-prime moduli. Berlekamp-Massey needs division, so it needs a prime modulus; matrix powering itself works with any modulus.
• Hidden non-linearity. Coefficients that depend on n, or products of earlier terms, do not fit. The brute-force comparison will catch this if you test enough values.
• Too few terms for Berlekamp-Massey. With fewer than 2k terms it can return a shorter, wrong recurrence that fits the data you gave it.

Choosing a method

For Fibonacci specifically, the 2-by-2 matrix collapses into fast-doubling identities that are faster still; see fast Fibonacci.

What to do next

1. Implement mat_mul and mat_pow yourself and test them against direct iteration for the first 30 terms.
2. Take a recurrence with a constant and a linear term, build its augmented matrix on paper, add the prefix-sum row, and verify with code.
3. Implement Kitamasa and confirm it matches the matrix method on a random order-10 recurrence.
4. Generate 40 terms of a tiling count by brute force, recover its recurrence with Berlekamp-Massey, and check it against held-back terms.
5. Rewrite one multi-query solution to use precomputed powers and vector products, and measure the speed-up.