[BOJ] 백준 1351번 무한 수열 (Python)

문제

무한 수열 A는 다음과 같다.

• A0 = 1
• Ai = A⌊i/P⌋ + A⌊i/Q⌋ (i ≥ 1)

N, P와 Q가 주어질 때, AN을 구하는 프로그램을 작성하시오.

입력

• 첫째 줄에 3개의 정수 N, P, Q가 주어진다.

출력

• 첫째 줄에 AN을 출력한다.

제한

• 0 ≤ N ≤ 10¹²
• 2 ≤ P, Q ≤ 10⁹

예제

입력

7 2 3

출력

7

힌트

• ⌊x⌋는 x를 넘지 않는 가장 큰 정수이다.

풀이

import sys

def dfs(num):
    global P, Q, arr
    if num < 1:
        return 1
    elif num in arr:
        return arr[num]
    arr[num] = dfs(num//P) + dfs(num//Q)
    return arr[num]


if __name__ == "__main__":
    N, P, Q = map(int, sys.stdin.readline().split())
    arr = {}

    print(dfs(N))

arr을 list()로 하면 메모리 초과. 리스트를 사용하지 않고 count 변수를 만든 후 dfs에서 num이 0일 때마다 count를 1씩 증가 하는 방식으로 하면 대부분 잘 되지만 히든 테스트케이스가 있는 듯 하다. 시간초과 발생. 따라서 딕셔너리를 사용해야 제대로 풀린다. 위 코드에서는 A0을 따로 지정하지 않고, 필요한 경우에 1을 리턴하는 것으로 대체하여 풀었다.