商汤 第三题 矩阵快速幂

主要根据公式计算出矩阵即可
笔试的时候并没调出来，不知道最后能过多少，测试用例过了，应该能过

添加 CPP 解法，这里需要改动的地方不多，如果有类似的题可以直接套用
按照评论区大佬给出的公式进行了修改，java 版本的没改，Java 大佬可以自行修改

CPP解法

#include <bits\stdc++.h>
using namespace std;

typedef long long ll;
typedef vector<ll> vec;
typedef vector<vec> mat;

const int MODE = 1e9 + 7;

mat multiply(mat& A, mat& B) {
    int m = A.size(), n = B[0].size();
    mat res(m, vec(n, 0));

    // 两个矩阵相乘的算法
    for(int i=0; i<m; i++)
        for(int j=0; j<n; j++)
            for(int k=0; k<A[i].size(); k++)
                res[i][j] = (res[i][j] + A[i][k] * B[k][j]) % MODE;

    return res;
}

mat pow(mat& A, long long n) {
    // 初始为单位矩阵
    mat res(A.size(), vec(A.size(), 0));
    for(int i=0; i<A.size(); i++)
        res[i][i] = 1;

    // 通过快速幂算法快速计算矩阵的 n 次方
    while(n > 0){
        if((n&1) == 1) res = multiply(A, res);
        n >>= 1;
        A = multiply(A, A);
    }

    return res;
}

long long getRes(int a, int b, int c, long long n, int f0) {
    // 计算得到的矩阵
    mat arr{{a, 1, 0, 0, 0, 0},
            {b, 0, 1, 0, 0, 0},
            {c, 0, 0, 0, 0, 0},
            {2, 0, 0, 1, 0, 0},
           {-1, 0, 0, 2, 1, 0},
        {32767, 0, 0, 1, 1, 1}};
    mat res = pow(arr, n);
    // 当 i 取 1 时的初始值
    mat F = {{f0, 0, 0, 1, 1, 1}};
    res = multiply(F, res);
    return res[0][0] % MODE;
}

int main(){
    ll n;
    int a, b, c, f0;
    cin >> n >> a >> b >> c >> f0;
    long long res = getRes(a, b, c, n, f0);
    cout << res << endl;
    return 0;
}

JAVA 解法

package com.interview.shoop;

import java.util.Scanner;

public class Main {
    final static int mod = 1000000007;

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        long n = sc.nextLong();
        int a = sc.nextInt();
        int b = sc.nextInt();
        int c = sc.nextInt();
        int f0 = sc.nextInt();
        System.out.println(helper(a, b, c, n, f0));
    }

    //矩阵乘法+快速幂
    public static int helper(int a, int b, int c, long n, int f0) {
        int[][] q = {
                {a, b, c, 1, 1, 32767},
                {1, 0, 0, 0, 0, 0},
                {0, 1, 0, 0, 0, 0},
                {0, 0, 0, 1, -4, 2},
                {0, 0, 0, 0, 1, -1},
                {0, 0, 0, 0, 0, 1}};
        int[][] res = pow(q, n);
        int[][] f = {{f0}, {0}, {0}, {2}, {-1}, {1}};
        res = multiply(res, f);       
        return res[0][0];
    }

    //矩阵快速幂
    public static int[][] pow(int[][] a, long n) {
        int[][] ret = {{1, 0, 0, 0, 0, 0}, {0, 1, 0, 0, 0, 0}, {0, 0, 1, 0, 0, 0}, {0, 0, 0, 1, 0, 0}, {0, 0, 0, 0, 1, 0}, {0, 0, 0, 0, 0, 1}};
        while (n > 0) {
            if ((n & 1) == 1) {
                ret = multiply(ret, a);
            }
            n >>= 1;
            a = multiply(a, a);
        }
        return ret;
    }

    //矩阵乘法
    public static int[][] multiply(int[][] a, int[][] b) {
        int m = a.length;
        int n = b[0].length;
        int[][] c = new int[m][n];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                int temp = 0;
                for (int k = 0; k < a[i].length; k++) {
                    temp = (temp + a[i][k] * b[k][j]) % mod;
                }
                c[i][j] = temp;
            }
        }
        return c;
    }
}