Skip to content
Embercode
✕
CombinatoricsLesson 8 of 11

45 min read

Linear recurrences and matrix exponentiation

Why this matters when you build with AI

A common contest sentence is "n can be as large as 10^18". A loop that runs n steps then needs 10^18 steps, which will not finish. The fix is to find a fixed set of numbers that describes the state after step i, and a matrix that turns the state at step i into the state at step i + 1. Raising that matrix to the power n then takes about 60 squarings for n near 10^18, instead of 10^18 loop steps.

AI tools know this idea, and they often write it wrong. The matrix is built with rows and columns swapped, or the power is one step too many. The sum of entries is taken over the wrong range, or the values are stored in int and overflow before the modulus is applied. None of these mistakes crashes the program, and the samples often pass, because samples use small n. The checks in this lesson catch each of these errors: a plain DP for small n, and a brute force for very small n.

The code is C++17. We compiled it with g++ -std=c++17 -O2 -Wall -Wextra using GCC 15.2.0 on Linux under WSL2. The output lines below come from those runs. The arithmetic does not depend on the machine, but the time a program takes does, and we say so where it matters.

Definitions

A linear recurrence is a rule that gives the next value from the last k values, with fixed numbers as multipliers. The Fibonacci numbers are the standard example: F(n) = F(n - 1) + F(n - 2), with F(0) = 0 and F(1) = 1. The order of the recurrence is k, the number of earlier values used. Fibonacci has order 2.

A state is the smallest set of numbers that lets you compute the next step. For Fibonacci the state after step n is the pair (F(n), F(n - 1)).

A transition matrix M is a k by k table such that the next state equals M times the current state. Matrix exponentiation means computing M^e, the matrix M multiplied by itself e times, by squaring, in about log2(e) steps.

A modulus is the number we take the remainder by. Here it is the prime 998244353. Every product and sum is reduced modulo this number, which keeps the numbers small. The first lesson of this module covers modular arithmetic and modular inverses.

Fibonacci as a matrix

Take the matrix

  • M = [[1, 1], [1, 0]]

Multiply it by the state (F(n), F(n - 1)). The first row gives F(n) + F(n - 1), which is F(n + 1). The second row gives F(n), which is the new second entry. So M turns the state at step n into the state at step n + 1.

After n steps the matrix M^n holds the Fibonacci numbers in a fixed pattern:

  • M^n = [[F(n + 1), F(n)], [F(n), F(n - 1)]]

Check with n = 1: M^1 = [[1, 1], [1, 0]] gives F(2) = 1, F(1) = 1, F(0) = 0. Check with n = 2: M^2 = [[2, 1], [1, 1]] gives F(3) = 2, F(2) = 1, F(1) = 1. Both match. So F(n) is the top-right entry of M^n. The program below prints F(50) mod 998244353 by this method, and by a plain loop, and the two values match.

A transition matrix for a local rule

The worked problem is a string count. Use the letters a, b and c. A string is valid when it has no "aa" and no "bc" as adjacent letters. Count the valid strings of length n.

The state is the last letter. Make a table T with T[x][y] = 1 when the pair "letter x then letter y" is allowed, and 0 otherwise. Reading the rules:

  • From a: "aa" is forbidden, so a can be followed by b or c. Row a is (0, 1, 1).
  • From b: "bc" is forbidden, so b can be followed by a or b. Row b is (1, 1, 0).
  • From c: nothing is forbidden. Row c is (1, 1, 1).

So T = [[0, 1, 1], [1, 1, 0], [1, 1, 1]].

Now count. A string of length 1 is any single letter, so there are 3. A string of length n is a sequence of n - 1 allowed steps, so it corresponds to a walk of n - 1 steps through T. The number of walks is the sum of all entries of T^(n - 1), because each walk starts at some letter and ends at some letter. That sum is the answer.

Check n = 2. The sum of all entries of T is 0 + 1 + 1 + 1 + 1 + 0 + 1 + 1 + 1 = 7. The 9 two-letter strings minus "aa" and "bc" leaves 7.

Check n = 3 by hand, without a matrix. Let r(x) be the sum of row x of T. Row sums are r(a) = 2, r(b) = 2, r(c) = 3. The number of valid strings of length 3 that start with x is the sum of r(y) over the allowed next letters y. Starting with a, the next letter is b or c, so 2 + 3 = 5. Starting with b, the next letter is a or b, so 2 + 2 = 4. Starting with c, all three letters are allowed, so 2 + 2 + 3 = 7. The total is 5 + 4 + 7 = 16.

The program below computes the answer by the matrix power and by a plain DP that keeps the count for each last letter. It also compares the two for every n from 1 to 300.

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using Mat = vector<vector<ll>>;
const ll MOD = 998244353;
 
Mat identity(int k) {
    Mat I(k, vector<ll>(k, 0));
    for (int i = 0; i < k; i++) I[i][i] = 1;
    return I;
}
 
// Multiply two k x k matrices modulo MOD. Cost is k^3 multiplications.
Mat mul(const Mat& A, const Mat& B) {
    int k = A.size();
    Mat C(k, vector<ll>(k, 0));
    for (int i = 0; i < k; i++)
        for (int t = 0; t < k; t++) {
            if (A[i][t] == 0) continue;
            for (int j = 0; j < k; j++) C[i][j] = (C[i][j] + A[i][t] * B[t][j]) % MOD;
        }
    return C;
}
 
// A^e by repeated squaring. About 2 log2(e) multiplications.
Mat mpow(Mat A, ll e) {
    Mat R = identity(A.size());
    while (e > 0) {
        if (e & 1) R = mul(R, A);
        A = mul(A, A);
        e >>= 1;
    }
    return R;
}
 
// Fibonacci through the matrix [[1,1],[1,0]]. F(n) is the top right entry of M^n.
ll fib_matrix(ll n) {
    Mat M = {{1, 1}, {1, 0}};
    return mpow(M, n)[0][1];
}
 
ll fib_dp(int n) {
    ll a = 0, b = 1;
    for (int i = 0; i < n; i++) {
        ll t = (a + b) % MOD;
        a = b;
        b = t;
    }
    return a;
}
 
// Count strings of length n over k letters where no adjacent pair (x, y) is in `bad`.
// Matrix T[x][y] = 1 when the pair "x then y" is allowed. The walks of n - 1 steps
// through T are the valid strings, so the answer is the sum of all entries of T^(n-1).
ll count_matrix(int k, const vector<pair<int, int>>& bad, ll n) {
    if (n == 0) return 1;
    Mat T(k, vector<ll>(k, 1));
    for (auto [x, y] : bad) T[x][y] = 0;
    Mat P = mpow(T, n - 1);
    ll total = 0;
    for (int i = 0; i < k; i++)
        for (int j = 0; j < k; j++) total = (total + P[i][j]) % MOD;
    return total;
}
 
// Plain DP over the length. dp[x] = valid strings of the current length that end in letter x.
ll count_dp(int k, const vector<pair<int, int>>& bad, int n) {
    vector<vector<bool>> forbidden(k, vector<bool>(k, false));
    for (auto [x, y] : bad) forbidden[x][y] = true;
    vector<ll> dp(k, 1);
    for (int len = 2; len <= n; len++) {
        vector<ll> nd(k, 0);
        for (int x = 0; x < k; x++)
            for (int y = 0; y < k; y++)
                if (!forbidden[x][y]) nd[y] = (nd[y] + dp[x]) % MOD;
        dp = nd;
    }
    ll total = 0;
    for (ll v : dp) total = (total + v) % MOD;
    return total;
}
 
int main() {
    cout << "Fibonacci by matrix: ";
    for (int n : {1, 2, 10, 50}) cout << "F(" << n << ")=" << fib_matrix(n) << " ";
    cout << "\n";
    cout << "F(50) by dp: " << fib_dp(50) << "\n";
 
    // Letters: 0 = a, 1 = b, 2 = c. Forbidden adjacent pairs: "aa" and "bc".
    vector<pair<int, int>> bad = {{0, 0}, {1, 2}};
    cout << "rule aa and bc forbidden\n";
    for (int n = 1; n <= 6; n++)
        cout << "n=" << n << " matrix " << count_matrix(3, bad, n) << " dp " << count_dp(3, bad, n) << "\n";
 
    bool ok = true;
    for (int n = 1; n <= 300; n++)
        if (count_matrix(3, bad, n) != count_dp(3, bad, n)) ok = false;
    cout << "matrix equals dp for n = 1..300: " << (ok ? "yes" : "no") << "\n";
    cout << "n = 10^18 by matrix: " << count_matrix(3, bad, 1000000000000000000LL) << "\n";
}

The output of this program, as we ran it:

Fibonacci by matrix: F(1)=1 F(2)=1 F(10)=55 F(50)=607336789
F(50) by dp: 607336789
rule aa and bc forbidden
n=1 matrix 3 dp 3
n=2 matrix 7 dp 7
n=3 matrix 16 dp 16
n=4 matrix 36 dp 36
n=5 matrix 81 dp 81
n=6 matrix 182 dp 182
matrix equals dp for n = 1..300: yes
n = 10^18 by matrix: 523503850

The n = 3 line matches the hand count of 16. The n = 10^18 line cannot be checked by the plain DP, which would need 10^18 steps. It is checked by the agreement up to n = 300, because the same matrix power code produces both. We also brute-forced every string up to length 7 with a short Python program, and the counts were 3, 7, 16, 36, 81, 182 and 409. The value at n = 7 is 409, and the matrix gives the same value.

The example uses the reduction mpow(T, n - 1), not mpow(T, n). That is the one-step error that the common mistakes section describes.

Fast power in O(k^3 log n)

The power is computed by repeated squaring. Write the exponent in binary. For example, 13 is 1101 in binary, which is 8 + 4 + 1. So M^13 = M^8 times M^4 times M. The loop in mpow reads one bit at a time. For each set bit it multiplies the result by the current square, and then it squares the current square.

For n up to 10^18, the exponent has at most 60 bits, since 2^60 is about 1.15 times 10^18. That gives at most 60 squarings and at most 60 extra multiplications. Each matrix product costs k^3 multiply-and-add steps for a k by k matrix. So the total is about 120 times k^3 steps.

For a 3 by 3 matrix this is about 3,000 steps, which is nothing. For a 10 by 10 matrix it is about 120,000. The number of steps is the figure to reason with. The time in seconds depends on your machine and your compiler, so measure it on your own machine before you decide on a limit.

A useful variant: if many queries use the same matrix and different exponents, compute the powers M, M^2, M^4, and so on up to M^(2^60) once. Each query then multiplies the state vector by the powers for its set bits. Each query costs k^2 times the number of bits, which is much less than k^3 times the number of bits.

Constants and modular multiplication

The following points change the speed and the correctness of the code.

Overflow in the products. The modulus 998244353 is less than 2^30. The largest product of two reduced values is (998244352)^2, which is about 9.96 times 10^17. This fits in a signed 64-bit integer, which holds up to about 9.2 times 10^18. A sum of 18 such products fits in an unsigned 64-bit integer, which holds up to about 1.84 times 10^19. A sum of 19 does not. The same limit holds for the modulus 10^9 + 7. So in the inner loop of a matrix product, you can add up to 18 products before you reduce, if the values are stored as unsigned 64-bit numbers. Our code reduces after each product, which is simpler and fast enough for small k.

Storing in 32 bits. If you store the values in int, the product of two of them is computed in int before the %, and it wraps around. Our demo program multiplies 999999999 by itself, and the output was:

32-bit product then mod: 808348673
64-bit product then mod: 64
exact product: 999999998000000001

The first line is wrong. The correct value of the product modulo 1000000007 is 64, as the second line shows. Widen one operand to long long or unsigned long long before multiplying.

Vector-of-vectors cost. vector<vector<ll>> allocates many small blocks of memory, and each access goes through two indirections. For a fixed small k, a plain array such as ll A[K][K] with K a constant is faster. Use it if the state size is known. Measure before you change it.

Exponent type. The exponent 10^18 fits in a signed 64-bit integer. Read it with long long, not int.

Identity start. The power loop starts from the identity matrix. If it starts from the zero matrix, the answer is 0 for every n. If it starts from the matrix M itself, the answer is off by one step.

When this does not apply

The method needs three things: the state has a fixed size k, the rule is the same at every step, and the rule is linear, which means each new value is a sum of old values with fixed coefficients. A recurrence such as f(n) = n times f(n - 1) has coefficients that depend on n, so it has no constant matrix. A counting problem whose state must remember the whole history of the string, or whose state grows with n, has no fixed k.

For a linear recurrence with constant coefficients and very large k, there are faster methods based on polynomial remainders, such as Kitamasa's algorithm. This lesson does not cover them.

Exercise: count strings under a local rule for n up to 10^18

Write a program for this problem. It is an original problem for this course.

A string is made of the letters 0, 1, 2 and 3. It is valid when no two adjacent letters are equal, and no 0 is directly followed by a 3. Given n, count the valid strings of length n, and print the count modulo 998244353.

Input: one integer n, with 1 at most n at most 10^18.

Output: one integer.

Examples from our runs, which a brute force over all 4^n strings confirmed for n up to 7:

  • n = 1 gives 4.
  • n = 2 gives 11.
  • n = 3 gives 30.
  • n = 4 gives 82.
  • n = 5 gives 224.
  • n = 6 gives 612.
  • n = 7 gives 1672.
  • n = 1000000000000000000 gives 910037098.

Hint: the state is the last letter, so the matrix is 4 by 4. Decide which pairs are forbidden, set those entries to 0 in an all-ones matrix, and use the sum of the entries of T^(n - 1). Check your matrix with n = 2 before you run a large n. The answer for n = 2 is the number of allowed pairs, which is 16 minus the 5 forbidden pairs, so 11.

Show the answer

The solution uses the functions in the worked example above, unchanged. Put the main below in place of the worked example's main, and keep everything above it. The forbidden pairs are the four equal pairs and the pair (0, 3). We compiled the combined file with g++ -std=c++17 -O2 -Wall -Wextra. It printed 4, 11, 30, 82 and 224 for n = 1 to 5, and 910037098 for n = 10^18. Its counts for n up to 300 agree with the plain DP, and the brute force agrees for n up to 7.

int main() {
    long long n;
    cin >> n;
    // Letters 0..3. Forbidden adjacent pairs: equal letters, and "0 then 3".
    vector<pair<int, int>> bad = {{0, 0}, {1, 1}, {2, 2}, {3, 3}, {0, 3}};
    cout << count_matrix(4, bad, n) << "\n";
}

Common mistakes

One step too many or too few. Using mpow(T, n) instead of mpow(T, n - 1) counts strings of length n + 1. For the worked rule it prints 7 for n = 1 instead of 3. Check the smallest case by hand every time.

Transposed matrix. If you build the matrix with rows and columns swapped, the sum of all entries of the power does not change, because the transpose of a power is the power of the transpose. The count is still right here. But the sum of row a of the power counts the strings that start with a, and the sum of column a counts the strings that end with a. These two numbers can differ, so the orientation matters once the question asks about a specific letter. Write the row as "from" and the column as "to", and write the start vector to match.

Overflow in the products. Storing values in int and multiplying them gives the wrong values as shown above. The symptom is a result that looks random and changes when you change the compiler flags, or a negative number.

Forgotten reduction. Without a reduction after each product, the entries grow exponentially, because the exponent doubles at each squaring. They overflow after a small number of squarings. The symptom appears only for large n, so the samples can pass.

A loop instead of a power. A loop over n steps is fine for small n and cannot finish for n near 10^18. The symptom is a time limit error, not a wrong answer.

Testing only one large value. An AI tool may give a matrix that prints a plausible number for n = 10^18 and is wrong for small n. Check every n from 1 to a few hundred against the plain DP, and check the smallest cases against a brute force.

A prompt to check AI output

Use this prompt on the program an AI tool gives you. Paste your code in the place of the marker.

Here is a C++ program that counts strings of length n under a local rule, using matrix
exponentiation, with n up to 10^18 and modulus 998244353. Check it for these points and name
the line for each:
1. Is the matrix power taken to n - 1, and why? Show the answer for n = 1 by hand.
2. Is each entry of the transition matrix set from the forbidden adjacent pairs, with rows as
   the first letter and columns as the second letter?
3. Can any product overflow a 64-bit integer before the % operator?
4. Does the power start from the identity matrix?
Then write a plain DP for the same rule, and a Python check that compares the two for every n
from 1 to 300. Print any mismatch.

A good answer names the line for each point, explains the n = 1 case by hand, and gives a test that you can run. An answer that only says "the matrix looks right" has not checked anything.

Checklist

  • Is the state the smallest set of values that determines the next step?
  • Is the transition matrix built from the rule for one adjacent pair?
  • Is the exponent n - 1 when the state is the last letter?
  • Does the power start from the identity matrix?
  • Are the values stored in 64-bit integers, and is every product reduced?
  • Is the exponent read as a 64-bit integer?
  • Did you compare the output with the plain DP for every n from 1 to a few hundred?
  • Did you check n = 1 and n = 2 by hand?

Summary

A linear recurrence with a fixed state becomes a matrix, and the answer after n steps is a matrix power. Repeated squaring computes the power in about log2(n) squarings. Most mistakes in this method are off-by-one exponents, transposed or mis-built matrices, and overflow in the products. The plain DP for small n is the fastest way to find them.

Practice

Each practice problem has a hint. Try it before you read the answer under the line "Show the answer".

  1. Count the strings of length 3 over the letters a and b that do not contain "aa". Hint: write the transition matrix for the two letters, then count the walks of length 2.

Show the answer There are 5. Listing them: aba, abb, bab, bba and bbb. The transition matrix is [[0, 1], [1, 1]], so the row sums are 1 for a and 2 for b. A string of length 3 has two steps. Starting from a, the only allowed next letter is b, whose row sum is 2, so that gives 2 strings. Starting from b, the next letters are a and b, with row sums 1 and 2, so that gives 3 strings. The total is 2 + 3 = 5.

  1. Find F(30) modulo 1000 with the Fibonacci matrix. Hint: F(30) is 832040, so check your answer against it.

Show the answer The answer is 40, because 832040 leaves 40 modulo 1000. The matrix M^30 has F(30) in its top-right entry, and reducing it modulo 1000 gives the same value.

  1. Two sequences are defined by a(1) = 1, b(1) = 0, a(n) = 2 a(n - 1) + b(n - 1), and b(n) = a(n - 1). Find a(4). Hint: write the state as the pair (a(n), b(n)) and find the 2 by 2 matrix.

Show the answer The matrix is [[2, 1], [1, 0]]. Stepping by hand: a(2) = 2 times 1 + 0 = 2, b(2) = 1. Then a(3) = 2 times 2 + 1 = 5, b(3) = 2. Then a(4) = 2 times 5 + 2 = 12.

  1. A program computes M^e for a 10 by 10 matrix and e = 10^18 with the plain repeated squaring in this lesson. Roughly how many multiply-and-add steps does it take? Hint: count the squarings and the extra multiplications, then multiply by k^3.

Show the answer There are about 60 squarings and at most 60 extra multiplications, so about 120 matrix products. Each product of 10 by 10 matrices takes 10^3 = 1000 multiply-and-add steps. The total is about 120,000 steps. This is a count of steps, not a time. The time depends on your machine.

  1. Why does the code in this lesson reduce after each product, and what would change if you summed up to 18 products before reducing? Hint: compare 18 times (998244352)^2 with 2^64.

Show the answer Each product of two reduced values is below about 9.96 times 10^17. Eighteen such products sum to about 1.79 times 10^19, which is below 2^64, about 1.84 times 10^19, so one reduction at the end is safe. Nineteen products sum to about 1.89 times 10^19, which is above 2^64, so the sum would wrap around. Reducing after each product is simpler, and the speed gain from waiting is small for small k.

Practice this topic: coding problems for Combinatorics →
← Counting paths on grids