이전에 백준 1629번 곱셉 문제를 풀면서 분할정복을 이용한 거듭제곱 문제에 익숙해져 있었기에 간단한 문제라고 생각하고 코드를 금방 짰다.
그러나 생각보다 고려해야할 점이 많은 까다로운 문제였다.
import sys
input = sys.stdin.readline
from collections import deque
N,B = map(int,input().split())
A = []
for i in range(N):
A.append(list(map(int,input().split())))
def calculate(A,B):
if B == 1:
return A
if B%2 == 0:
return multifly(calculate(A,B//2),calculate(A,B//2))
else:
return multifly(multifly(calculate(A,B//2),calculate(A,B//2)),A)
def multifly(A,B):
result = [[0]*N for i in range(N)]
for i in range(N):
for j in range(N):
value = 0
for k in range(N):
value += A[i][k] * B[k][j] % 1000
result[i][j] = value%1000
return result
sol = calculate(A,B)
for i in sol:
print(" ".join(map(str,i)))처음 제출한 코드는 이러하다.
그러나 이 코드는 시간초과로 통과되지 못하였다.
그러다 질문 게시판에서 메모리제이션이 필요할 것 같다는 조언을 보고 새로 코드를 짰다.
import sys
input = sys.stdin.readline
from collections import deque
N,B = map(int,input().split())
A = []
memory = {}
def remainder(x):
return x%1000
for i in range(N):
A.append(list(map(remainder,map(int,input().split()))))
def calculate(A,B):
if B == 1:
return A
if memory.get(B):
return memory[B]
if B%2 == 0:
memory[B] = multifly(calculate(A,B//2),calculate(A,B//2))
return memory[B]
else:
memory[B] = multifly(multifly(calculate(A,B//2),calculate(A,B//2)),A)
return memory[B]
def multifly(A,B):
result = [[0]*N for i in range(N)]
for i in range(N):
for j in range(N):
value = 0
for k in range(N):
value += A[i][k] * B[k][j] % 1000
result[i][j] = value%1000
return result
sol = calculate(A,B)
for i in sol:
print(" ".join(map(str,i)))행렬의 거듭제곱값을 memory 딕셔너리에 저장하고, 딕셔너리에 값이 있으면 딕셔너리 값을 리턴하고, 없으면 딕셔너리에 계산해서 채워넣는 방식이다.
참고로 이 코드는 처음에는 통과되지 못하였는데, 치명적인 실수를 했기 때문이었다.
바로 multifly 함수에만 나머지 계산을 해주고 행렬 자체에는 나머지계산을 하지 않아서 B가 1이고 A에 1000을 넘어가는 값이 있으면 나머지계산이 되지 않았던 것이다.
나와 같은 실수를 한 사람들이 많은지 이것도 질문게시판에서 찾아 해결할 수 있엇다.
그러다 다른 사람들의 코드를 보고, 굳이 메모리 dictionary가 필요없다는 사실을 알게 되었는데, 그냥 함수 자체에서 B//2 계산값을 저장해주면 되는 것이었다!
최종 코드는 이러하다.
import sys
input = sys.stdin.readline
from collections import deque
N,B = map(int,input().split())
A = []
def remainder(x):
return x%1000
for i in range(N):
A.append(list(map(remainder,map(int,input().split()))))
def calculate(A,B):
if B == 1:
return A
cal2 = calculate(A,B//2)
if B%2 == 0:
return multifly(cal2,cal2)
else:
return multifly(multifly(cal2,cal2),A)
def multifly(A,B):
result = [[0]*N for i in range(N)]
for i in range(N):
for j in range(N):
value = 0
for k in range(N):
value += A[i][k] * B[k][j] % 1000
result[i][j] = value%1000
return result
sol = calculate(A,B)
for i in sol:
print(" ".join(map(str,i)))
참고로 최종 코드는 64ms, 2차 코드는 60ms만에 통과하였는데, 아무래도 memory dictionary를 만드는 방식이 복잡하긴 해도 약간이라도 시간 단축을 시킬 수 있는 코드인 것 같다.
오늘의 교훈) 문제의 조건을 꼼꼼이 따지고 실수하지 말자. 또 함수는 최대한 적게 시현하여 시간복잡도를 줄이자.