[Algorithm] 백준 - 11444번 피보나치 수 6 (Swift)

문제 소개

피보나치를 구하는 방법은 재귀, dp등 다양하지만 이 문제는 n의 범위가 10^18 인 무시무시한 문제다. 당연 O(n)으로 문제를 푼다면 틀리게 된다.

문제 풀이

이 문제는 좀더 수학적으로 접근해야 한다. 우선 너무나도 유명한 피보나치 수열의 일반항은 아래와 같다.

이것을 행렬로 나타낸다면 아래와 같이 나타낼 수 있다.

오른쪽을 쭈우우우욱 풀어준다면... 아래와 같이 결국 표현된다. 이제 우리는 행렬의 n-1 제곱만 구해주면 된다!

2x2 행렬의 곱은 아래와 같이 구할 수 있다.

func multiply(_ a: [[Int]], _ b: [[Int]]) -> [[Int]] {
    let x = (a[0][0] * b[0][0] + a[0][1] * b[1][0]) % mod
    let y = (a[0][0] * b[0][1] + a[0][1] * b[1][1]) % mod
    let z = (a[1][0] * b[0][0] + a[1][1] * b[1][0]) % mod
    let w = (a[1][0] * b[0][1] + a[1][1] * b[1][1]) % mod
    return [[x, y], [z, w]]
}

거듭제곱은 아래와 같이 0일때 항등행렬, 절반으로 나눠 거듭제곱하는 방식으로 시간복잡도를 줄여준다.

func power(_ base: [[Int]], _ exp: Int) -> [[Int]] {
    if exp == 0 { return [[1, 0], [0, 1]] }
    if exp == 1 { return base }
    if exp % 2 == 0 {
        let half = power(base, exp / 2)
        return multiply(half, half)
    } else {
        return multiply(base, power(base, exp - 1))
    }
}

만든 함수를 바탕으로 f(n)을 아래와 같이 쉽게 구할 수 있다!

func fib(_ n: Int) -> Int {
    if n == 0 { return 0 }
    if n == 1 { return 1 }
    let base = [[1, 1], [1, 0]]
    let result = power(base, n - 1)
    return result[0][0]
}

정답 코드

import Foundation

let n = Int(readLine()!)!
let mod = 1_000_000_007

func multiply(_ a: [[Int]], _ b: [[Int]]) -> [[Int]] {
    let x = (a[0][0] * b[0][0] + a[0][1] * b[1][0]) % mod
    let y = (a[0][0] * b[0][1] + a[0][1] * b[1][1]) % mod
    let z = (a[1][0] * b[0][0] + a[1][1] * b[1][0]) % mod
    let w = (a[1][0] * b[0][1] + a[1][1] * b[1][1]) % mod
    return [[x, y], [z, w]]
}

func power(_ base: [[Int]], _ exp: Int) -> [[Int]] {
    if exp == 0 { return [[1, 0], [0, 1]] }
    if exp == 1 { return base }
    if exp % 2 == 0 {
        let half = power(base, exp / 2)
        return multiply(half, half)
    } else {
        return multiply(base, power(base, exp - 1))
    }
}

func fib(_ n: Int) -> Int {
    if n == 0 { return 0 }
    if n == 1 { return 1 }
    let base = [[1, 1], [1, 0]]
    let result = power(base, n - 1)
    return result[0][0]
}

print(fib(n))

아주 만족스런 결과가 나왔다 🥸