【leetcode 1425. 带限制的子序列和】【矩阵幂快速运算】

class Solution {
    public int countVowelPermutation(int n) {
        long[][] matrix = new long[][]{
                {0, 1, 1, 0, 1},
                {1, 0, 1, 0, 0},
                {0, 1, 0, 1, 0},
                {0, 0, 1, 0, 0},
                {0, 0, 1, 1, 0}
        };
        int mod = 1_000_000_007;
        matrix = pow(matrix, n - 1, mod);
        long ans = 0l;
        for (int i = 0; i < 5; i++) {
            for (int j = 0; j < 5; j++) {
                ans += matrix[i][j];
                ans %= mod;
            }
        }
        return (int) ans;
    }

    public long[][] pow(long[][] matrix, int p, int mod) {
        int n = matrix.length;
        int m = matrix[0].length;
        if( p == 0){
            long[][] ans = new long[n][m];
            for(int i=0;i<n;i++){
                for(int j = 0;j<m;j++){
                    if(i == j){
                        ans[i][j] = 1;
                    }
                }
            }
            return ans;
        }
        if (p == 1) {
            long[][] ans = new long[n][m];
            for(int i=0;i<n;i++){
                for(int j = 0;j<m;j++){
                    ans[i][j] = matrix[i][j];
                }
            }
            return ans;
        }

        int h = n;

        long[][] newmatrix = pow(matrix, p / 2, mod);
        long[][] ans = mutiply(mod, newmatrix, newmatrix);
        if (p % 2 == 1) {
            ans = mutiply(mod, ans, matrix);
        }

        return ans;
    }


    private long[][] mutiply(int mod, long[][] matrix1, long[][] matrix2) {
        int n = matrix1.length;
        int m = matrix1[0].length;
        int h = matrix2[0].length;

        long[][] ans = new long[n][h];
        for (int i = 0; i < n; i++) {
            for (int k = 0; k < h; k++) {
                long t = 0l;
                for (int j = 0; j < m; j++) {
                    long c = matrix1[i][j] * matrix2[j][k];
                    c %= mod;
                    t += c;
                    t %= mod;
                }
                ans[i][k] = t;
            }
        }
        return ans;
    }
}