백준 10961번: School canteen (Matrix)

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

10961번: School canteen

---

아이디어가 기존에 포스팅한 문제와 똑같아서 운 좋게 풀이 시간을 크게 단축할 수 있었습니다.

---

우선 조건에 맞게 $N$일동안 식사를 제공하는 경우의 수를 $A_{N}$이라 하고,

 

$B_{N}(i,j)$를 $N-j+1$일부터 $N$일까지 $i$번째 메뉴를 제공한 경우의 수라고 하겠습니다.

 

이때 $A_{N}$과 $B_{N}$의 관계가 다음과 같습니다.

$$ A_{N} = \sum_{1 \leq i \leq k} \sum_{j=1}^{l-1} B_{N}(i,j) $$

그리고 $B_{N}$의 점화 관계를 살펴보겠습니다.

 

첫번째로

$$ B_{N}(i,1) = A_{N-1} - \sum_{j=1}^{l-1} B_{N-1}(i,j) $$

임을 알 수 있습니다. $i$번째 메뉴가 $N-1$일 째에 제공될 수 없기 때문입니다.

 

두번째로

$$ B_{N}(i,k) = B_{N-1}(i,k-1) \;\; (1 < k < l) $$

임을 알 수 있습니다. $i$번째 메뉴가 $N-1$일 째에 $k-1$번 연속으로 제공되어야 하는 차례이기 때문입니다.

 

이를 바탕으로, 앞선 포스팅과 같은 전개 방식으로, 다음을 얻을 수 있습니다.

$$ \sum_{j=1}^{l-1} B_{N}(i,j) = A_{N-1} - A_{N-l} + \sum_{j=1}^{l-1} B_{N-l}(i,j) $$

$$ \therefore A_{N} = kA_{N-1} - (k-1)A_{N-l} $$

---

여기서 행렬 거듭제곱을 사용하여 시간복잡도 $O(l^{3} \times \log N)$을 얻을 수 있습니다.

$$ \begin{bmatrix} A_{N} \\ A_{N-1} \\ A_{N-2} \\ \vdots \\ A_{N-l+1} \end{bmatrix} = \begin{bmatrix} k & 0 & 0 & \cdots & 0 & 1-k \\ 1 & 0 & 0 & \cdots & 0 & 0 \\ 0 & 1 & 0 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & \cdots & 1 & 0 \end{bmatrix}^{N-l} \begin{bmatrix} A_{l} \\ A_{l-1} \\ A_{l-2} \\ \vdots \\ A_{1} \end{bmatrix} $$

$l \leq 2 \times 10^{2}$ 이지만, 시간제한이 5초이므로 시간은 차고 넘칩니다. (실제로 커팅을 잘 하면 400ms 미만으로 줄어들었습니다.)

 

초항의 경우, 자명하게 $x < l$일 때 $A_{x} = k^{x}$입니다.

 

$A_{l}$ 역시 쉽게 $A_{l}=k^{l}-k$임을 알 수 있습니다. 다만 $A_{0}$의 경우 이 문제가 요구하는 정답은 1인데, 행렬 거듭제곱이 아닌 kitamasa를 쓸 때는 이슈가 생깁니다.

 

행렬 거듭제곱으로 구현 시 주의할 점으로는, $2^{31} < 4 \times 10^{9} + 9 < 2^{32}$라는 것과 음수에 대한 모듈러 처리가 있겠습니다.

 

다음은 제 코드입니다.

#include <stdio.h>
#define mod 4000000009

typedef struct {
    unsigned long long array[200][200];
} Matrix;

Matrix matrix_multiply_modular (Matrix A, Matrix B, int size);
Matrix matrix_power_modular (Matrix A, int size, long long P);

int main() {
    
    long long N, L, K;
    scanf("%lld %lld %lld", &N, &L, &K);
    
    int size = L;
    
    Matrix simple = {{{}}};
    simple.array[size - 1][0] = K;
    for (int i = size - 2; i >= 0; i--) {
        simple.array[i][0] = simple.array[i + 1][0] * K % mod;
    }
    simple.array[0][0] = (simple.array[0][0] + mod - K) % mod;
    
    if (N <= L) {
        printf("%lld", N == 0 ? 1 : simple.array[L - N][0]);
        return 0;
    }
    
    Matrix base = {{{}}};
    for (int i = 0; i + 1 < size; i++) {
        base.array[i + 1][i] = 1;
    }
    base.array[0][0] = K;
    base.array[0][size - 1] = 1 + mod - K;
    
    base = matrix_multiply_modular(matrix_power_modular(base, size, N - L), simple, size);
    printf("%llu", base.array[0][0]);
    
    return 0;
}

Matrix matrix_multiply_modular (Matrix A, Matrix B, int size) {
    Matrix result = {{{}}};
    for (int i = 0; i < size; i++) {
        for (int j = 0; j < size; j++) {
            __int128 temp = 0;
            for (int k = 0; k < size; k++) {
                temp += A.array[i][k] * B.array[k][j];
            }
            result.array[i][j] = temp % mod;
        }
    }
    return result;
}

Matrix matrix_power_modular (Matrix A, int size, long long P) {
    Matrix result = {{{}}};
    for (int i = 0; i < size; i++) {
        result.array[i][i] = 1;
    }
    while (P) {
        if (P % 2) {
            result = matrix_multiply_modular(result, A, size);
        }
        A = matrix_multiply_modular(A, A, size);
        P /= 2;
    }
    return result;
}

(C11, 376ms, 3808KB, 제출번호 67114882)

---

이 문제 역시 kitamasa법을 활용하여 실행시간을 크게 단축시킬 수 있습니다. (4ms, 제출번호 67113732)

 

그 부분은 따로 포스팅하도록 하겠습니다.

 

행렬 거듭제곱을 구현할 때 내부를 128 bit로 늘리는 건 좋지 않았습니다.

 

의외로 놀란 건, i-j-k loop에다가 size 인수까지 제공하지 않는, 커팅을 전혀 거치지 않은 구현도 시간을 넉넉히 통과한다는 것이었습니다. (1696ms, 제출번호 67115085)

 

한편으로 글 서두에 27308번을 언급했는데, 그 문제의 notation을 가져오면 $X=Y$이고 $Z$의 범위가 크게 줄어들었다고 볼 수 있겠습니다.