10830번: 행렬 제곱

10830번: 행렬 제곱

사용 언어: C++

문제 요약

크기가 N*N인 행렬 A가 주어진다. 이때, A의 B제곱을 구하는 프로그램을 작성하시오. 수가 매우 커질 수 있으니, A^B의 각 원소를 1,000으로 나눈 나머지를 출력한다.

첫째 줄에 행렬의 크기 N과 B가 주어진다. (2 ≤ N ≤  5, 1 ≤ B ≤ 100,000,000,000)

둘째 줄부터 N개의 줄에 행렬의 각 원소가 주어진다. 행렬의 각 원소는 1,000보다 작거나 같은 자연수 또는 0이다.

첫째 줄부터 N개의 줄에 걸쳐 행렬 A를 B제곱한 결과를 출력한다.

풀이

행렬 제곱을 단순하게 구현하면 A를 B번 곱할 수 있는데 B는 최대 10^11이므로 불가능하다.

그래서 분할 정복 거듭제곱을 적용하면 되는데,

B가 짝수 지수일 때

A^B = (A^(B/2)) × (A^(B/2))

그리고 B가 홀수 지수일 때

A^B = (A^(B-1)) × A

이 과정을 반복하면 지수가 절반씩 줄어든다. 따라서 O(logB) 단계만 필요해서 총 시간 복잡도가 O(N^3 * logB)가 된다. (N은 현재 문제에서 최대 5로 작음)

행렬의 제곱은 단순하게 각 행렬끼리 곱하고 1000으로 나눈 나머지를 저장하면 되고, 분할 정복 제곱은 아래 코드로 설명해보면

Matrix power(Matrix a, long long exp, int n) {
    // 행렬 a를 exp번 거듭제곱한 결과를 반환하는 함수
    // n은 행렬의 크기 (n x n)

    if (exp == 1) {
        // 재귀 탈출 조건
        // A^1 = A 그대로 반환
        for (int i = 0; i < n; i++)
            for (int j = 0; j < n; j++)
                a[i][j] %= MOD;
        return a;
    }

    // 1. 절반 제곱 구하기
    // A^exp를 바로 계산하는 대신, A^(exp/2)를 먼저 구한다
    Matrix half = power(a, exp / 2, n);

    // 2. 절반 제곱을 자기 자신과 곱해서 (A^(exp/2))^2을 만든다
    // 즉, A^(exp/2) × A^(exp/2) = A^exp (exp가 짝수일 경우)
    Matrix result = multiply(half, half, n);

    // 3. exp가 홀수라면 A를 한 번 더 곱해줘야 한다
    // 왜냐하면 exp = 2k+1 형태일 때, A^exp = (A^k)^2 × A 이기 때문이다
    if (exp % 2 == 1) 
        result = multiply(result, a, n);

    return result;
}

 

이에 따라 최종 코드는

답

#include <iostream>
#include <vector>

using namespace std;

const int MOD = 1000;
using Matrix = vector<vector<long long>>;

Matrix multiply(const Matrix& a, const Matrix &b, int n)
{
    Matrix result(n, vector<long long>(n, 0));
    
    for(int i=0; i<n; i++)
    {
        for(int k=0; k<n; k++)
        {
            for(int j=0; j<n; j++)
            {
                result[i][j] += (a[i][k] * b[k][j]);
            }
        }
        
        for(int j=0; j<n; j++)
        {
            result[i][j] %= MOD;
        }
    }
    
    return result;
}

// 분할 정복 제곱
Matrix power(Matrix a, long long exp, int n)
{
    if(exp == 1)
    {
        Matrix result = a;
        for(int i=0; i<n; i++)
        {
            for(int j=0; j<n; j++)
            {
                result[i][j] %= MOD;
            }
        }
        return result;
    }
    
    Matrix half = power(a, exp/2, n);
    Matrix result = multiply(half, half, n);
    if(exp % 2 == 1)
    {
        result = multiply(result, a, n);
    }
    
    return result;
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    
    int n;
    long long b;
    cin >> n >> b;
    
    Matrix A(n, vector<long long>(n));
    for(int i=0; i<n; i++)
    {
        for(int j=0; j<n; j++)
        {
            cin >> A[i][j];
        }
    }
    
    Matrix result = power(A, b, n);
    
    for(int i=0; i<n; i++)
    {
        for(int j=0; j<n; j++)
        {
            cout << result[i][j] % MOD << " ";
        }
        cout << '\n';
    }
    
    return 0;
}