[Baekjoon]백준 10830 행렬 제곱(골드 4) - Python, C++

문제설명

예제

문제를 살펴보면 이전에 풀이하였던 곱셈에서 제곱으로 레벨이 조금 올라간 모습이다

이전 문제를 풀어오면서 분할 정복에 대한 알고리즘은 조금 익혔기 때문에 간단하게만 설명을 하자면

그냥 일반적인 방법으로 제곱을 하게 되면 O(n)의 행렬 곱셈을 필요로 하지만, 분할 정복을 사용하면

시간 복잡도를 O(logN)으로 줄일 수 있다 n = 1인 경우 행렬을 return하면 되고,

짝수인 경우는 A^n = A^n/2 x A^n/2, 홀수인 경우는 A^n = A^n//2 x A^n//2 x A와 같은 방식으로 곱하면 된다

def mul(a, b): # 행렬 곱 구하는 함수
    res = [[0] * size for _ in range(size)]
    for i in range(size):
        for j in range(size):
            for k in range(size):
                res[i][j] += a[i][k] * b[k][j] % 1000
    return res

def sol(x, n): # 행렬 거듭제곱 구하는 함수
    if n == 1:
        return x
    tmp = sol(x, n // 2)
    if n % 2: # n이 홀수인 경우
        return mul(mul(tmp, tmp), x)
    else: # n이 짝수인 경우 
        return mul(tmp, tmp)
    
A = []
size, pow_num = map(int, input().split())
for _ in range(size):
    A.append(list(map(int, input().split())))
    
result = sol(A, pow_num)
for i in range(size): 
    for j in range(size):
        result[i][j] = result[i][j] % 1000 # 문제 조건
        
for i in result:
    print(*i)

파이썬 코드로 작성하면 다음과 같다 행렬 곱을 구한 후에 1000으로 나누어주면서(문제 조건) 계산할 값을 줄여주고,

최종적으로 반환받은 값도 1000으로 나누어준 나머지를 저장해주면 된다

C++로 작성을 할 때 꽤나 오랜 시간이 걸렸는데, 로직 그대로 C++로 작성하였더니 틀렸습니다 라는 문구가 등장하였다

...어 뭐지? 싶었고 계속해서 시도해보았다 int로 선언한 부분을 long long으로 바꾸어도 보았고, 일부 로직을 수정해보아도

틀렸습니다는 그대로 발생해서 결국 지피티의 도움을 조금 받았다 이전의 코드를 보여주자면

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

vector<vector<long long>> mul(const vector<vector<long long>> &a, const vector<vector<long long>> &b, int size) {
    vector<vector<long long>> res(size, vector<long long>(size, 0));
    for (int i = 0; i < size; i++) {
        for (int j = 0; j < size; j++) {
            for (int k = 0; k < size; k++) {
                res[i][j] += a[i][k] * b[k][j];
                res[i][j] %= 1000;  // 문제 조건
            }
        }
    }
    return res;
}

vector<vector<long long>> sol(const vector<vector<long long>> &x, int n, int size) {
    if (n == 1)
        return x;
    vector<vector<long long>> tmp = sol(x, n / 2, size);
    if (n % 2 == 1)  // n이 홀수인 경우
        return mul(mul(tmp, tmp, size), x, size);
    else  // n이 짝수인 경우
        return mul(tmp, tmp, size);
}

int main() {
    int size;
    long long pow_num;
    cin >> size >> pow_num;
    vector<vector<long long>> A(size, vector<long long>(size));
    for (int i = 0; i < size; i++) {
        for (int j = 0; j < size; j++) {
            cin >> A[i][j];
        }
    }

    vector<vector<long long>> result = sol(A, pow_num, size);
    for (int i = 0; i < size; i++) {
        for (int j = 0; j < size; j++) {
            result[i][j] %= 1000;  // 문제 조건
        }
    }

    for (const auto& row : result) {
        for (const auto& elem : row) {
            cout << elem << ' ';
        }
        cout << endl;
    }

    return 0;
}

저 코드에서 요약에 적힌 부분을 수정하여 주었고, 그 덕분에 드디어 맞았습니다! 를 볼 수 있었다

파이썬과 C++의 동작 과정의 차이를 몸소 느낄 수 있었던 시간이었다

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

vector<vector<long long>> mul(const vector<vector<long long>> &a, const vector<vector<long long>> &b, int size) {
    vector<vector<long long>> res(size, vector<long long>(size, 0));
    for (int i = 0; i < size; i++) {
        for (int j = 0; j < size; j++) {
            for (int k = 0; k < size; k++) {
                res[i][j] += a[i][k] * b[k][j];
                res[i][j] %= 1000;  // 문제 조건
            }
        }
    }
    return res;
}

vector<vector<long long>> sol(const vector<vector<long long>> &x, long long n, int size) {
    if (n == 1) {
        vector<vector<long long>> res = x;
        for (int i = 0; i < size; i++)
            for (int j = 0; j < size; j++)
                res[i][j] %= 1000;  // 1000으로 나누기
        return res;
    }

    vector<vector<long long>> tmp = sol(x, n / 2, size);
    tmp = mul(tmp, tmp, size);
    
    if (n % 2 == 1)  // n이 홀수인 경우
        return mul(tmp, x, size);
    return tmp;
}

int main() {
    int size;
    long long pow_num;
    cin >> size >> pow_num;
    vector<vector<long long>> A(size, vector<long long>(size));
    for (int i = 0; i < size; i++) {
        for (int j = 0; j < size; j++) {
            cin >> A[i][j];
        }
    }

    vector<vector<long long>> result = sol(A, pow_num, size);

    for (const auto& row : result) {
        for (const auto& elem : row) {
            cout << elem << ' ';
        }
        cout << endl;
    }

    return 0;
}

수많은 시도의 흔적...