[CodeUp/2753] 수열의 n번째 항 (Python)

Problem : 수열의 n번째 항 (codeup.kr)

Status : Solved

Time : 00:19:39

 

---

 

문제 설명

 

KDS는 점화식에 관심이 많다.

그래서 프로그래밍을 이용해서 어떤 점화식이 주어졌을 때, 그 수열의 n번째항을 구하고 싶다.

이제 점화식에 대한 정보가 주어질 때, n번째항을 구하는 프로그램을 짜서 KDS를 도와주자.

구하고자 하는 수열의 점화식이

$$ F_n=m_1∗F_{n−1}+m_2∗F_{n−2} $$

위와 같을 때, 제 0항 a,제 1항 b, 그리고 점화식에서의 계수 m1, m2이 공백으로 주어진다. 그리고 n이 주어진다.

이 때 n번째 항의 값을 출력하되 1,000,000,007로 나눈 나머지를 출력하시오.

 

입력 및 출력

 

입력

제 0항 a, 제 1항 b, 그리고 선형점화식에서의 계수 m1, m2와 자연수 n이 공백으로 입력된다.

(단, n은 100,000,000,000보다 같거나 작다. a,b,m1,m2는 100이하의 자연수. )

 

출력

n번째 항을 출력하되, 1,000,000,007로 나눈 나머지를 출력한다.

 

입력 예시   

0 1 2 1 2

 

출력 예시

2

 

---

 

풀이

 

피보나치 수열 및 선형점화식 문제는 시간복잡도 O(N)으로 푸는 방법을 쉽게 떠올릴 수 있다. 하지만 본 문제에서는 최대 100000000000의 케이스를 요구하므로, O(N)보다 작은 시간복잡도의 풀이를 생각해야 한다.

 

이에 대해 백준 블로그에서 자세히 다룬 바 있다.

피보나치 수를 구하는 여러가지 방법 (acmicpc.net)

피보나치 수를 구하는 여러가지 방법

 

핵심은 행렬. 선형점화식 문제를 행렬로 표현한다면 행렬의 거듭제곱 꼴로 나타낼 수 있고, 거듭제곱을 구하는 시간복잡도는 O(logN)으로 적용 가능하다.

위 문제를 행렬식으로 일반화한다고 생각하자. 먼저 다음과 같은 두 식을 가정하겠다.

 

$$ F_n = m_1 * F_{n-1} + m_2 * F_{n-2} $$

$$ F_{n-1} = F_{n-1} $$

 

이를 행렬로 표현하면...

 

$$ \begin{pmatrix}
 F_n \\
 F_{n-1}
\end{pmatrix} = 
\begin{pmatrix}
  m_1 & m_2 \\
  1 & 0 
\end{pmatrix}
\begin{pmatrix}
 F_{n-1} \\
 F_{n-2}
\end{pmatrix} $$

 

이 된다. 눈치챘을지도 모르지만, 이는 모든 n값에 대해 적용하므로 이를 풀어 쓸 수 있다. 이를 정리하면 

 

$$ \begin{pmatrix}
 F_n \\
 F_{n-1}
\end{pmatrix} = 
\begin{pmatrix}
  m_1 & m_2 \\
  1 & 0 
\end{pmatrix}^{n-1}
\begin{pmatrix}
 F_{1} \\
 F_{0}
\end{pmatrix} $$

 

가 되므로, 위 행렬의 거듭제곱을 계산하면 쉽게 구할 수 있다. 행렬의 거듭제곱은 다음 코드로 갈음하도록 하겠다.

 

풀이 코드

MOD = 1000000007

def matmul(mat1, mat2) :
  new_mat = [[0]*2 for _ in range(2)]
  for i in range(2) :
    for j in range(2) :
      tmp = 0
      for k in range(2):
        tmp += mat1[i][k]*mat2[k][j]
      new_mat[i][j] = tmp % MOD
  return new_mat

def mat_pow(mat, p) :
  if p == 1 :
    return mat

  pmat = mat_pow(mat, p//2)
  pmat = matmul(pmat, pmat)
  if p % 2 == 0 :
    return pmat
  else :
    return matmul(pmat, mat)

def solution():
  a, b, m1, m2, n = map(int, input().split())
  if n == 1 :
    print(b)
    return

  mat = [[m1, m2], [1, 0]]
  pmat = mat_pow(mat, n-1)
  print((b*pmat[0][0] + a*pmat[0][1])%MOD)

solution()

풀이 완료!