백준 6673번: Firepersons

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

6673번: Firepersons

---

지문 하단에 구해야 하는 값이 명시되어 있는데, $i$가 주어지면

$$ a_{i} = \sum_{j=1}^{k} a_{i-j}b_{j} $$

을 구하면 됩니다. 이 문제에서 쿼리 개수가 정해져있지 않고, $i$가 최대 $10^{9}$ 이므로 단순 반복문으로는 힘들겠습니다. 행렬 제곱을 활용하면, $i > k$에 대해

$$ \begin{bmatrix} a_{i} \\ a_{i-1} \\ a_{i-2} \\ \vdots \\ a_{i-k+1} \end{bmatrix} = \begin{bmatrix} b_{1} & b_{2} & \cdots & b_{k-1} & b_{k} \\ 1 & 0 & \cdots & 0 & 0 \\ 0 & 1 & \cdots & 0 & 0 \\ \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & \cdots & 1 & 0 \end{bmatrix} \begin{bmatrix} a_{i-1} \\ a_{i-2} \\ a_{i-3} \\ \vdots \\ a_{i-k} \end{bmatrix} $$

이므로 binary exponentiation을 이용해 $a_{i}$를 쿼리 당 $O(k^{3}\log i)$에 구할 수 있습니다.

 

물론 쿼리 개수가 충분히 많거나, $k$가 $10^{3}$ 이상으로 커지는 경우라면 kitamasa 등 알고리즘으로 복잡도를 떨궈야 합니다. 이 문제에서는 다행히 여기까지만 해도 충분했습니다.

 

제 코드는 다음과 같습니다.

#include <stdio.h>
#define mod 10000

typedef struct {
	long long array[100][100];
} Matrix;

Matrix matrix_multiply_modular (Matrix A, Matrix B, int S);
Matrix matrix_power_modular (Matrix A, int P, int S);

int main() {
	
	while(1) {
		int k;
		scanf("%d", &k);
		if (k == 0) {
			return 0;
		}
		
		int a[100] = {};
		for (int idx = 0; idx < k; idx++) {
			scanf("%d", &a[idx]);
		}
		
		int b[100] = {};
		for (int idx = 0; idx < k; idx++) {
			scanf("%d", &b[idx]);
		}
		
		int i;
		scanf("%d", &i);
		if (i < k) {
			printf("%d\n", a[i]);
			continue;
		}
		
		Matrix Base = {{{}}};
		for (int idx = 0; idx < k; idx++) {
			Base.array[0][idx] = b[idx];
			Base.array[idx + 1][idx] = 1;
		}
		Base = matrix_power_modular(Base, i-k+1, k);
		
		int result = 0;
		for (int idx = 0; idx < k; idx++) {
			result = (result + Base.array[0][idx] * a[k-1-idx]) % mod;
		}
		printf("%d\n", result);
	}
	return 0;
}

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

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

(C11, 52ms, 1620KB, 제출번호 65339745)

---

제한시간이 1초이므로 약간의 커팅은 필요합니다.

 

커팅을 전혀 하지 않았을 때는 시간을 초과합니다. 예를 들어 modulo를 상수로 놓지 않으면 시간을 초과합니다.

 

modulo를 상수로 놓고, 3중 loop를 ikj로 했을 때는 무려 976ms가 나왔습니다. (제출번호 65339859)

 

여기서 ijk로 순서를 바꾼 뒤 % 연산 횟수를 최적화했을 때 440ms가 나왔습니다. (제출번호 65339823)

 

대부분의 쿼리가 $k$ 값이 작을 거라고 기대하면, 행렬 연산 함수에 행렬의 사이즈를 같이 줄 수 있습니다. 여기까지 커팅한 것이 본문 코드입니다.