[백준 10830 - C++] 행렬 제곱 : 분할 정복

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

10830번: 행렬 제곱

문제

크기가 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제곱한 결과를 출력한다.

 

풀이

위와 같은 성질을 이용해서 재귀 호출을 하는 코드를 짜면 된다.
알고리즘에 대한 아이디어는 매우 쉽게 떠올릴 수 있는 문제이다.
여기서 중요한 것은 사소한 조건들을 잘 고려해야한다는것이다ㅠㅠ
우선 b의 범위가 int를 넘어간다. 따라서 long long 형을 이용한다!

또한 나는 재귀함수의 매개변수로는 n과 제곱해야할 횟수를 전달하고 함수의 return값을 2차원 배열로 해주었다.
그렇기 때문에 재귀 실행 중에 메모리 누수가 많이 발생할 수 있어서 메모리 초과를 막기 위한 처리를 해주어야한다.
따라서 동적 할당 해준 배열들을 제때 제때 적절한 위치에서 free해주는 것이 중요하다!
그러다보니 코드가 조금 많이 더러워졌다... 가독성도 좀 많이 떨어진다.
하나의 함수로 해결하려 하지 말고 함수를 추가로 더 사용했으면 더 깔끔한 코드가 되었을 텐데 하는 생각이든다.

#include <iostream>
#include <cstdio>
#include <vector>
#include <stdlib.h>

using namespace std;
int mat[6][6];

int** sol(int n, long long k){
    int **use;
    int **temp;
    int **result;
    long long new_k = k/2;
    
    result = (int **)calloc(n, sizeof(int*));
    for (int i = 0; i < n; i++)
        result[i] = (int *)calloc(n, sizeof(int));
    
    temp = (int **)calloc(n, sizeof(int*));
    for (int i = 0; i < n; i++)
        temp[i] = (int *)calloc(n, sizeof(int));
    
    if(k==2){
        for(int i=0; i<n; i++){
            for(int j=0; j<n; j++){
                for(int k=0; k< n; k++)
                    result[i][j] += mat[i][k]*mat[k][j];
                result[i][j] = result[i][j]%1000;
            }
        }
        return result;
    }
    
    else if(k==3){
        use = sol(n, 2);
        for(int i=0; i<n; i++){
            for(int j=0; j<n; j++){
                for(int k=0; k< n; k++)
                    result[i][j] += use[i][k]*mat[k][j];
                temp[i][j] = result[i][j]%1000;
            }
        }
    }
    else{
        use = sol(n, new_k);
        
        for(int i=0; i<n; i++){
            for(int j=0; j<n; j++){
                for(int k=0; k< n; k++)
                    result[i][j] += use[i][k]*use[k][j];
                result[i][j] = result[i][j]%1000;
            }
        }
        
        if(k%2){
            for(int i=0; i<n; i++){
                for(int j=0; j<n; j++){
                    for(int k=0; k< n; k++)
                        temp[i][j] += result[i][k]*mat[k][j];
                    temp[i][j] = temp[i][j]%1000;
                }
            }
        }
        
    }
    
    if(k%2 == 0){
        for(int i=0; i<n; i++){
            free(use[i]);
            free(temp[i]);
        }
        free(use);
        free(temp);
        
        return result;
    }
    else{
        for(int i=0; i<n; i++){
            free(use[i]);
            free(result[i]);
        }
        free(use);
        free(result);
        return temp;
    }
}

int main(void){
    int n;
    long long b;
    int **result;
    
    scanf("%d %lld", &n, &b);
    for(int i=0; i<n; i++){
        for(int j=0; j<n; j++){
            scanf("%d", &mat[i][j]);
        }
    }
    
    if(b==1){
        for(int i=0; i<n; i++){
            for(int j=0; j<n; j++){
                printf("%d ", mat[i][j]%1000);
            }
            printf("\n");
        }
        return 0;
    }
    
    result = sol(n, b);
    
    for(int i=0; i<n; i++){
        for(int j=0; j<n; j++){
            printf("%d ", result[i][j]);
        }
        printf("\n");
    }
}