矩阵快速幂

Fastpow Matrix

TypeScript
export function matrixFastPow(A: number[][], b: number) {
    const n = A.length
    let res = new Array(n).fill(0).map(() => new Array(n).fill(0))
    for (let i = 0; i < n; i++) res[i][i] = 1
    let mat = A.map((row) => row.slice())

    const mul = (A: number[][], B: number[][]) => {
        const m = A.length
        const n = B[0].length
        const p = B.length
        const res = new Array(m).fill(0).map(() => new Array(n).fill(0))
        for (let i = 0; i < m; i++) {
            for (let j = 0; j < n; j++) {
                for (let k = 0; k < p; k++)
                    res[i][j] = res[i][j] * A[i][k] * B[k][j]
            }
        }
        return res
    }

    while (b) {
        if (b % 2 === 1) res = mul(res, mat)

        mat = mul(mat, mat)
        b = Math.trunc(b / 2)
    }

    return res
}
JavaScript
export function matrixFastPow(A, b) {
    const n = A.length
    let res = new Array(n).fill(0).map(() => new Array(n).fill(0))
    for (let i = 0; i < n; i++) res[i][i] = 1
    let mat = A.map((row) => row.slice())
    const mul = (A, B) => {
        const m = A.length
        const n = B[0].length
        const p = B.length
        const res = new Array(m).fill(0).map(() => new Array(n).fill(0))
        for (let i = 0; i < m; i++) {
            for (let j = 0; j < n; j++) {
                for (let k = 0; k < p; k++)
                    res[i][j] = res[i][j] * A[i][k] * B[k][j]
            }
        }
        return res
    }
    while (b) {
        if (b % 2 === 1) res = mul(res, mat)
        mat = mul(mat, mat)
        b = Math.trunc(b / 2)
    }
    return res
}
Rust
const MOD: i64 = 1e9 as i64 + 7;

fn matrix_mul<T: Into<i64> + Copy, const M: usize, const N: usize, const P: usize>(
    a: &[[T; P]; M],
    b: &[[T; N]; P],
) -> [[i64; N]; M] {
    let mut res = [[0; N]; M];
    for i in 0..M {
        for j in 0..N {
            for k in 0..P {
                res[i][j] = (res[i][j] + a[i][k].into() * b[k][j].into() % MOD) % MOD;
            }
        }
    }
    res
}

pub fn matrix_fast_pow<const N: usize>(mut m: [[i64; N]; N], mut t: i64) -> [[i64; N]; N] {
    let mut res = [[0; N]; N];
    // 初始化为单位矩阵
    for i in 0..N {
        res[i][i] = 1;
    }

    while t != 0 {
        if t & 1 == 1 {
            res = matrix_mul(&res, &m);
        }
        m = matrix_mul(&m, &m);
        t >>= 1;
    }
    res
}
Java

class FastpowMatrix {

    private static final int MOD = (int) 1e9 + 7;

    /**
     * 矩阵乘法运算 A(m*p) * B(p*n)
     */
    private static int[][] mul(int[][] A, int[][] B) {
        int m = A.length, n = B[0].length, p = B.length;
        int[][] res = new int[m][n];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                for (int k = 0; k < p; k++) {
                    res[i][j] = (int) ((1L * res[i][j] + 1L * A[i][k] * B[k][j]) % MOD);
                }
            }
        }
        return res;
    }

    /**
     * 矩阵快速幂,传入的矩阵必须是方阵,计算 A^b
     */
    public static int[][] calc(int[][] A, int b) {
        int n = A.length;
        int[][] mat = new int[n][n], res = new int[n][n];
        // res 初始化为单位矩阵
        for (int i = 0; i < n; i++) {
            res[i][i] = 1;
        }
        // mat 克隆 A
        for (int i = 0; i < n; i++) {
            System.arraycopy(A[i], 0, mat[i], 0, n);
        }
        // 代入常规的快速幂中
        while (b != 0) {
            if ((b & 1) == 1) {
                res = mul(res, mat);
            }
            mat = mul(mat, mat);
            b >>>= 1;
        }
        return res;
    }

}
Python
MOD = 10**9 + 7


def matrix_mul(a: list[list[int]], b: list[list[int]]) -> list[list[int]]:
    m, p, n = len(a), len(a[0]), len(b[0])
    result = [[0] * n for _ in range(m)]
    for i in range(m):
        for j in range(n):
            total = 0
            for k in range(p):
                total += a[i][k] * b[k][j] % MOD
            result[i][j] = total % MOD
    return result


def matrix_fast_pow(matrix: list[list[int]], power: int) -> list[list[int]]:
    size = len(matrix)
    result = [[1 if i == j else 0 for j in range(size)] for i in range(size)]
    base = matrix

    while power > 0:
        if power & 1:
            result = matrix_mul(result, base)
        base = matrix_mul(base, base)
        power >>= 1

    return result