[C++] 행렬 제곱 (백준 10830번)

https://www.acmicpc.net/problem/10830

 

#include <iostream>
#include <vector>
using namespace std;

int n;
long long b;
int mod = 1000;

vector<vector<int>> mpow(vector<vector<int>> a, long long c) {
    vector<vector<int>> res(n, vector<int>(n));

    if (c == 1) {
        return a;
    } else if (c % 2 == 0) {
        vector<vector<int>> tmp = mpow(a, c / 2);
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                int rest = 0;
                for (int k = 0; k < n; k++) {
                    rest += tmp[i][k] * tmp[k][j];
                    rest %= mod;
                }
                res[i][j] = rest;
            }
        }
    } else {
        vector<vector<int>> tmp = mpow(a, c - 1);
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                int rest = 0;
                for (int k = 0; k < n; k++) {
                    rest += tmp[i][k] * a[k][j];
                    rest %= mod;
                }
                res[i][j] = rest;
            }
        }
    }

    return res;
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    cin >> n >> b;

    vector<vector<int>> a(n, vector<int>(n));
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            cin >> a[i][j];
            a[i][j] %= mod;
        }
    }

    vector<vector<int>> ans = mpow(a, b);

    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            cout << ans[i][j] << ' ';
        }
        cout << '\n';
    }
}

b가 짝수면, b를 절반으로 줄여서 얻은 행렬끼리 행렬 곱셈을 진행한다.

b가 홀수면, b에서 1을 빼서 구한 행렬과 원래 행렬을 곱한다. 

b가 1이면 그대로 반환하여 종료 조건을 만들어둔다.

여기서 주의할 점은 입력이 처음부터 행렬의 원소로 1000이 들어오고 b는 1일 수도 있다. 그렇기 때문에 입력 받을 때 한번 나머지 연산을 해주어야 한다. (AC)