[C++] 백준 5095 - Matrix Powers [Gold4]

출처

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

문제가 영어로 되어 있어서 당황할 수는 있지만, 단순히 문제에서 주어진 예시만 보더라도 알 수 있듯이, N,M,P를 입력 받아서 NxN 행렬을 만든 뒤 각각의 성분을 M으로 나눈 나머지를 P번 제곱하면 된다.
다만, 주의할 점은 출력 형식이다. 예제가 1개밖에 없어서 0 0 0을 입력해서 프로그램을 종료시키는 것이라고 생각해서 0 0 0을 새로운 변수로 선언해 입력값을 주었지만, 그것이 아니라 0 0 0도 N,M,P이고 N,M,P가 모두 0이 입력될 때까지 무한 반복으로 Testcase를 입력 받아서 행렬 제곱 연산을 수행하는 것이였다. 또한 각 Testcase마다 한 줄씩 띄워져 있어야한다.
이 부분 때문에 3번이나 출력 형식 오류가 났었다.

풀이 전략

행렬의 제곱연산을 하기 위해서는 행렬의 곱셈 방식과 그 과정에서 index의 변화를 이해해야 한다. 또한 그냥 연산을 하게 될 경우 값은 제대로 출력이 되지만 시간초과 문제가 발생할 수 있기 때문에 알고리즘 분류에 적혀있는 대로 분할정복을 이용하여 풀이하여야 한다.

풀이

그림이 좀 대충 한 것 처럼 보이긴 하지만, 설명을 하자면..

• 1행 1열 : 빨간색 형광펜으로 그은 일직선끼리의 연산 = A[0][0]*B[0][0]+A[0][1]*B[1][0]
• 1행 2열 : 파란색 형광펜으로 그은 일직선끼리의 연산 = A[0][0]*B[0][1]+A[0][1]*B[1][1]

이 둘을 비교하면 현재 1행에 대한 연산이고 앞에 있는 행렬이 A라고 하면 A의 행은 1행(index = 0)으로 고정되어 있다는 것을 알 수 있다. 그리고 뒤에 있는 행렬을 B라고 하면 B는 열이 고정되어 있는 것을 알 수 있다.
따라서 이 둘의 관계를 통해 다음과 같은 연산 수식을 유도할 수 있다.

    for(int i=0; i<N; i++){
        for(int j=0; j<N; j++){
            for(int k=0; k<N; k++){
                Matrix[i][j] += A[i][k]*B[k][j];
            }
        }
    }

그 다음으로 분할정복을 사용해야 하는데, 지수가 짝수일 때와 홀수일 때를 나누어서 지수가 짝수일 경우는 제곱 연산을 하고 지수를 나누고 홀수인 경우는 행렬의 모든 원소가 1인 행렬을 만들어서 그 행렬과 곱해준다.

vector<vector<int>> MatPow(vector<vector<int>> A, int exp, int ModNum){
    int size = A.size();
    vector<vector<int>> result(size,vector<int>(size, 0));

    for(int i=0; i<size; i++){
        result[i][i] = 1; // 단위행렬로 초기화
    }

    while(exp>0){
        if(exp%2==1){
            result = MatMul(result, A, ModNum); // 홀수일때는 단위행렬에 A를 곱해준다.
        }
        A = MatMul(A, A, ModNum); // A를 제곱해준다.
        exp/=2;
    }
    return result;
}

소스 코드

1번의 시간초과와 2번의 출력형식 오류가 있었다. 그 오류의 원인들을 분석하고 정답 코드를 분석해보려고 한다.

1. 시간 초과가 발생한 코드

/*문제 : https://www.acmicpc.net/problem/5095
  알고리즘 : 수학, 선형대수
  티어 : Gold4
*/
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int N,M,P;
vector<vector<int>> Mat;

vector<vector<int>> MatMul(vector<vector<int>> A, vector<vector<int>> B, int ModNum){
    vector<vector<int>> Temp(N,vector<int>(N, 0));
    for(int i=0; i<N; i++){
        for(int j=0; j<N; j++){
            for(int k=0; k<N; k++){
                Temp[i][j] += A[i][k]*B[k][j];
                Temp[i][j] %= ModNum;
            }
        }
    }
    return Temp;
}

int main(void){
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    
    while(1){
        cin>>N>>M>>P;

        Mat.resize(N,vector<int>(N));
        for(int i=0; i<N; i++){
            for(int j=0; j<N; j++){
                cin>>Mat[i][j];
            }
        }

        for(int i=1; i<P; i++){
            Mat = MatMul(Mat,Mat,M);
        }
        for(int i=0; i<N; i++){
            for(int j=0; j<N; j++){
                cout<<Mat[i][j]<<" ";
            }
            cout<<"\n";
        }
        Mat.clear();
    }
    return 0;
}

• 분석을 해보면, 이미 MatMul 함수에서 3중 반복문으로 인해 시간복잡도가 O(N^3)인데 main에서 P번 제곱을 해야 해서 한번 더 반복문을 돌고 있기 때문에 O(N^3∗P)의 시간 복잡도가 발생해서 시간 초과가 발생한 것 같다.

2. 출력 형식 오류

for(int i=0; i<N; i++){
   for(int j=0; j<N; j++){
      cout<<Mat[i][j]<<" ";
   }
   cout<<"\n";

   if(i<N-1) 
   	  cout<<"\n";
}

• 문제에서 한 칸을 뛰라고 해서 i<N-1 이 코드를 넣었는데, 이렇게 하게 될 경우 행이 끝날 때 마다 또 한칸 씩 띄워지는 문제가 있어서 저 코드를 밖으로 보냈더니 조금 찝찝하지만 정답처리가 되었다.

3. 정답 코드

/*문제 : https://www.acmicpc.net/problem/5095
  알고리즘 : 수학, 선형대수
  티어 : Gold4
*/
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int N,M,P;
vector<vector<int>> Mat;

vector<vector<int>> MatMul(vector<vector<int>> A, vector<vector<int>> B, int ModNum){
    vector<vector<int>> Temp(N,vector<int>(N, 0));
    for(int i=0; i<N; i++){
        for(int j=0; j<N; j++){
            for(int k=0; k<N; k++){
                Temp[i][j] += A[i][k]*B[k][j];
                Temp[i][j] %= ModNum;
            }
        }
    }
    return Temp;
}

vector<vector<int>> MatPow(vector<vector<int>> A, int exp, int ModNum){
    int size = A.size();
    vector<vector<int>> result(size,vector<int>(size, 0));

    for(int i=0; i<size; i++){
        result[i][i] = 1; // 단위행렬로 초기화
    }

    while(exp>0){
        if(exp%2==1){
            result = MatMul(result, A, ModNum); // 홀수일때는 단위행렬에 A를 곱해준다.
        }
        A = MatMul(A, A, ModNum); // A를 제곱해준다.
        exp/=2;
    }
    return result;
}

int main(void){
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    
    while(1){
        cin>>N>>M>>P;
        if (N==0 && M==0 && P==0){
            break;
        }

        Mat.resize(N,vector<int>(N));
        for(int i=0; i<N; i++){
            for(int j=0; j<N; j++){
                cin>>Mat[i][j];
            }
        }

        Mat = MatPow(Mat,P,M);

        for(int i=0; i<N; i++){
            for(int j=0; j<N; j++){
                cout<<Mat[i][j]<<" ";
            }
            cout<<"\n";
        }
        Mat.clear(); // 다음 테스트케이스를 위해 초기화
        cout<<'\n'; // 이 곳으로 옮김
    }
    return 0;
}

출력은 제대로 되지만, 찝찝했던 이유는 예를 들어 testcase가 2개인 경우 t1 출력 (한 줄 띄고) t2 출력 까지는 문제가 없지만, t2가 출력된 이후에도 한 줄이 더 띄워지기 때문이다.

위의 사진이 내가 test로 돌려본 내용인데, 문제에서 첫 번째 행렬만 예시로 줘서 두번째 3x3행렬은 임의로 1 2 3을 조합해서 만들어보았다.
0 0 0을 입력한 이후 제대로 출력이 되지만, 마지막에 또 한줄이 띄워지는 것을 확인할 수 있다.
문제에서는 testcase끼리의 구분을 위해 한 줄을 띄우라고 한 것 같은데, 그래서 이미 구분은 되었기 때문에 마지막 한 줄이 띄워지는 것이 문제가 되지 않는 것이라고 생각을 해본다.

참고

유사한 문제

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