Matrix Exponentiation, in depth: building transition matrices, a worked example by hand, overflow-safe code, walks, min-plus shortest paths and automata

One step as a matrix

Take the recurrence a(n) = a(n-1) + 2 a(n-2). If you keep a state vector holding the last two values, v(n) = [a(n), a(n-1)], then one step is a linear map: the new first entry is a(n) + 2 a(n-1), and the new second entry is just the old first entry. Written as a matrix, v(n+1) = M v(n) with M = [[1, 2], [1, 0]]. Apply it n times and you get v(n+1) = M^n v(1).

Matrix multiplication is associative, so M^n can be built exactly like a number power: square M to get M^2, M^4, M^8 and so on, and multiply together the powers that correspond to the 1 bits of n. For n = 10^18 that is about 60 squarings and at most 60 extra multiplications, rather than 10^18 steps. With a k-by-k matrix each multiplication costs k^3 operations, so the total is O(k^3 log n). That is the whole algorithm; the skill is in building M.

Building the transition matrix

Building the transition matrix follows a recipe. Choose a state vector that contains everything the next step needs, then write each new state entry as a linear combination of old ones. Each row of M is one such combination.

• A recurrence of order k. For a(n) = c1 a(n-1) + ... + ck a(n-k), the state is the last k values. The first row of M holds c1 ... ck, and the rows below form a shifted identity that moves each value down one slot.
• A constant term. For a(n) = a(n-1) + 2 a(n-2) + 3, add a constant 1 to the state. Its row is [0, ..., 0, 1] so it stays 1, and the first row gets a 3 in that column.
• A polynomial term in n. For + n, keep both n and 1 in the state; n advances by the row that adds the 1 to it.
• A running sum. To get S(n) = a(0) + ... + a(n), add S to the state with the row S(n+1) = S(n) + a(n+1), expanding a(n+1) through its own row.
• Several coupled sequences. Put all of them in the state; each gets a row describing how it updates from all of them.

Two conventions cause most bugs, so fix them once. This page uses column vectors and multiplies on the left, v(n+1) = M v(n), so row i of M produces entry i of the new state. Row-vector write-ups transpose M; never mix the two. And decide which index the starting vector holds, because the exponent is the number of steps from that index, not n itself.

Worked example by hand

Compute a(10) for a(n) = a(n-1) + 2 a(n-2) + 3, with a(0) = 1 and a(1) = 2. The state is v(n) = [a(n), a(n-1), 1] and

        | 1  2  3 |                       | 2 |
    M = | 1  0  0 |      v(1) = [a(1), a(0), 1] = | 1 |
        | 0  0  1 |                       | 1 |

We need v(10) = M^9 v(1). In binary 9 is 1001, so M^9 = M^8 times M. Squaring gives:

M^2 = [[3, 2, 6],     M^4 = [[11, 10, 30],    M^8 = [[171, 170, 510],
       [1, 2, 3],            [ 5,  6, 15],           [ 85,  86, 255],
       [0, 0, 1]]            [ 0,  0,  1]]           [  0,   0,   1]]

M^9 = M^8 M = [[341, 342, 1023],
               [171, 170,  510],
               [  0,   0,    1]]

The first entry of M^9 v(1) is 341 times 2 plus 342 times 1 plus 1023 times 1, which is 682 + 342 + 1023 = 2047. Check by stepping: the sequence is 1, 2, 7, 14, 31, 62, 127, 254, 511, 1022, 2047, so a(10) = 2047. The second entry, 171 times 2 plus 170 plus 510 = 1022, is a(9), as the state promises. Three squarings and one extra multiplication replaced nine steps; for n = 10^18 the same code does about 120 multiplications of 3-by-3 matrices.

Implementation and overflow

The implementation is two functions. Python integers do not overflow, but reduce modulo a prime anyway, because problems ask for answers modulo 10^9 + 7 and unbounded integers become slow as they grow.

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 k in range(m):
            a = Ai[k]
            if a == 0:
                continue                      # sparse rows are common; skip them
            Bk = B[k]
            for j in range(p):
                Ci[j] = (Ci[j] + a * Bk[j]) % mod
    return C

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

def a(n):                                     # the worked example
    if n < 2:
        return [1, 2][n]
    T = [[1, 2, 3], [1, 0, 0], [0, 0, 1]]
    return mat_mul(mat_pow(T, n - 1), [[2], [1], [1]])[0][0]

The loop order i, k, j reads rows of B contiguously, which matters in compiled languages. In C, C++, Java or Rust, overflow is the trap. With a modulus near 10^9, each product is near 10^18, close to the signed 64-bit limit of about 9.2 times 10^18, so you cannot sum k products and reduce once at the end. Reduce after every product, or accumulate in a 128-bit integer, or in an unsigned 64-bit accumulator reduce at least every 18 products. Negative coefficients must be stored as mod minus the value, because the remainder of a negative number differs between languages.

// C++: reduce after every multiply-add; values stay below MOD squared plus MOD.
for (int i = 0; i < n; ++i)
  for (int k = 0; k < n; ++k) {
    uint64_t a = A[i][k];
    if (!a) continue;
    for (int j = 0; j < n; ++j)
      C[i][j] = (C[i][j] + a * B[k][j]) % MOD;
  }

Walks in graphs and other semirings

Let A be the adjacency matrix of a directed graph, with A[i][j] = 1 when there is an edge from i to j. Then the (i, j) entry of A^k is the number of walks of exactly k edges from i to j. The proof is the definition of multiplication: (A^2)[i][j] sums A[i][m] A[m][j] over every middle vertex m, which counts the two-edge walks through each m, and induction does the rest.

Example: three vertices with edges 0 to 1, 0 to 2, 1 to 2, 2 to 0 and 2 to 1. Computing A^4 gives [[2, 3, 3], [2, 2, 1], [1, 3, 4]]. The top-left 2 says there are two closed walks of length 4 from vertex 0, and listing them confirms it: 0, 2, 0, 2, 0 and 0, 2, 1, 2, 0. A Markov chain after k steps is the same computation with a stochastic matrix, in floating point.

The algorithm works for any operations that form a semiring: an addition that is associative and commutative, a multiplication that is associative and distributes over it, and identities for both. Replace addition with min and multiplication with plus, and the product computes shortest paths: (W^k)[i][j] is the length of the shortest walk from i to j using exactly k edges. The min-plus identity has 0 on the diagonal and infinity everywhere else, and a missing edge is infinity, not 0.

INF = float("inf")

def minplus_mul(A, B):
    n = len(A)
    return [[min(A[i][k] + B[k][j] for k in range(n)) for j in range(n)] for i in range(n)]

# edges: 0->1 (4), 0->2 (1), 1->2 (2), 2->0 (3), 2->1 (1)
W = [[INF, 4, 1], [INF, INF, 2], [3, 1, INF]]
W2 = minplus_mul(W, W)   # [[4, 2, 6], [5, 3, inf], [inf, 7, 3]]
W3 = minplus_mul(W2, W)  # [[9, 7, 4], [inf, 9, 5], [6, 4, 9]]

W2[0][1] = 2 is the path 0, 2, 1. W3[0][1] = 7 is the cheapest walk with exactly three edges, 0, 1, 2, 1, which is longer than the two-edge path because the edge count is fixed. To get at most k edges instead, put 0 on the diagonal of W so a walk can stay put. Raising to a power by squaring then answers questions like the cheapest route with exactly a billion hops in O(n^3 log k), something Dijkstra cannot express because it does not track edge counts.

Counting strings with an automaton

Counting strings that avoid a pattern is a walk count on an automaton. Build the KMP automaton of the pattern, whose state is the length of the longest prefix of the pattern that is a suffix of the text read so far, as described in the KMP algorithm. Drop the state that means the pattern has been completed. The number of length-n strings that avoid the pattern is the number of length-n walks from state 0 that stay inside the remaining states.

For strings over {a, b} that avoid aba, the states are 0, 1 (read a) and 2 (read ab). From 0, a goes to 1 and b stays at 0; from 1, a stays at 1 and b goes to 2; from 2, b goes to 0 and a would complete the pattern, so it is dropped. With column-vector convention, M[to][from] counts transitions: M = [[1, 0, 1], [1, 1, 0], [0, 1, 0]]. Summing the column for state 0 of M^n gives 7, 12, 21 and 351 for n = 3, 4, 5 and 10, which matches a brute-force enumeration: of the 8 strings of length 3 only aba itself is excluded. The same construction with an Aho-Corasick automaton handles several forbidden patterns at once.

Cost and how to go faster

The cost is O(k^3 log n) for a k-state problem, so the state size dominates. A 3-by-3 matrix is trivial; a 200-state automaton needs 8 million multiply-adds per product and about 60 products for n = 10^18, which is fine in C++ and slow in pure Python. Ways to do better:

• Many queries with the same matrix. Precompute M, M^2, M^4 and so on once. Then answer each query by multiplying the start vector by the needed powers: each step is a vector-matrix product, O(k^2), so a query costs O(k^2 log n) instead of O(k^3 log n).
• Only one entry of a linear recurrence. The Kitamasa method uses the characteristic polynomial (Cayley-Hamilton) to compute x^n modulo that polynomial, costing O(k^2 log n), or less with fast polynomial multiplication.
• A recurrence you do not know. If you can compute the first 2k terms by brute force and the sequence satisfies some linear recurrence of order at most k, Berlekamp-Massey finds it, and then you can use either method.

To check the asymptotic claims and constant factors in your own setting, the analysis tools in Big-O, in depth apply directly.

Common mistakes

Always test against a brute-force stepper for small n before trusting the fast version: the bugs above all produce plausible numbers.

What to do next

1. Implement mat_mul and mat_pow from this page and reproduce a(10) = 2047 and the M^2, M^4, M^8 matrices by hand.
2. Add a brute-force stepper and compare it with the fast version for n from 0 to 50.
3. Practice the construction: add a running sum to the example, then a term in n, and check both against brute force.
4. Count walks of length k in a small graph and list them by hand for k = 3 to confirm one entry.
5. Switch to min-plus, compute shortest paths with exactly k edges, and add zeros on the diagonal to get at most k edges.
6. Build the KMP automaton for a pattern, count avoiding strings, and verify with enumeration for n up to 12.
7. Port the multiplication to a 64-bit language and test with values near the modulus to catch overflow.
8. When the state grows past a few dozen, try precomputed powers for many queries, or Kitamasa for one term of a long recurrence.