[99클럽 코테스터디 25일차 TIL / 백준] 1351 - 무한 수열 (DP, HashMap)

99클럽 코테스터디 25일차 TIL
KeyWord : DP, Map

문제

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

• $A_0 = 1$
• $A_i = A_{⌊i/P⌋} + A_{⌊i/Q⌋} (i ≥ 1)$

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

입력

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

출력

첫째 줄에 AN을 출력한다.

제한

• $0 ≤ N ≤ 10^{12}$
• $2 ≤ P, Q ≤ 10^9$

 

---

 

풀이

매우 간단한 DP 문제이다. 처음부터 문제에 수열의 점화식도 주어졌으므로, 그대로 메모제이션을 하면서 코드를 구현하면 될 듯하다. (bottom-up 방식이 아닌, top-down의 재귀 방식으로 구현하면 쉽다.)

하지만 이 정도만 고려하고 곧장 코드를 구현한다면 반드시 실패할 것인데, 아래의 조건 때문이다.

• $0 ≤ N ≤ 10^{12}$

N의 범위가 32bit 정수 자료형을 초과한다. 즉, int가 아닌 long 자료형을 사용해야 한다는 것이다.

이제 통과할 수 있을까?

하지만 대부분은 메모리 초과로 실패한다. 아마 대부분의 분들은 아래와 같이 메모제이션을 하려고 했을 것이기 때문이다.

Long[] dp = new Long[N+1];

다른 DP 문제였다면 문제가 없었겠지만, 이번에도 N의 범위가 문제다.

여기서 N의 값으로 $10^{12}$이 들어온다면? $10^{12}$ 크기의 Long 배열이 만들어진다. $8bytes \times 10^{12}$라는 엄청난 양의 메모리를 차지하게 된다.

즉, 여기서 당연하게도 메모리 초과가 발생한다.

 

점화식을 살펴보자.

$A_i = A_{⌊i/P⌋} + A_{⌊i/Q⌋} (i ≥ 1)$

조금 더 이해하기 쉽게 코드로 표현하면 아래와 같다.

dp[i] = dp[i/P] + dp[i/Q]

여기서 $2 ≤ P, Q$ 이므로 인덱스는 매번 점화식이 진행될 때마다 절반식 줄어든다. 즉, 그냥 N 크기의 배열을 이용하면 낭비되는 공간이 매우 많다는 것을 알 수 있다.

• N = 100
• P = 2
• Q = 2

이렇게 입력이 주어졌을 때, 실제로 사용되는 인덱스는 100, 50, 25, 12, 6, 3, 1, 0이다. 위에서 설명한한 것처럼 대부분의 인덱스는 사용되지 않는다는 것을 알 수 있다.

핵심은 이렇게 공간이 낭비되지 않게끔 하는 것이다. 어떤 방법이 있을까?

 

이것을 해결하려면 기존의 dp 문제를 해결할 때의 습관에서 벗어나야 한다. 생각해보면 메모제이션을 할 때 반드시 배열을 사용할 필요는 없다. 그냥 이 인덱스일 때는 이 값을 가진다, 라는 것만 기록할 수 있으면 된다.

인덱스와 값. key와 value. 즉, Map 자료 구조가 떠오르지 않는가?

Map<Long, Long> dp = new HashMap<>();  

이렇게 메모제이션을 위한 해쉬 맵을 선언한 후, key에는 인덱스, value에는 메모제이션을 위한 값을 저장하면 된다.

값을 꺼낼 때의 시간 복잡도 또한 $O(1)$이므로 배열을 이용하는 것과 별 차이가 없다.

이제 진짜로 코드로 구현만 하면 된다.

 

코드

import java.io.BufferedReader;  
import java.io.IOException;  
import java.io.InputStreamReader;  
import java.util.HashMap;  
import java.util.Map;  
import java.util.StringTokenizer;  

public class Main {  
    static Map<Long, Long> dp = new HashMap<>();  

    public static void main(String[] args) throws IOException {  
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));  
        StringTokenizer stk = new StringTokenizer(br.readLine());  

        Long n = Long.parseLong(stk.nextToken());  
        Long p = Long.parseLong(stk.nextToken());  
        Long q = Long.parseLong(stk.nextToken());  


        // N 크기의 DP 배열 선언은 0 <= N <= 10^12 조건 때문에 불가능  
        // HashMap<idx, value>로 해결  
        dp.put(0L, 1L);  

        // top-down 방식  
        solution(n, p, q);  
        System.out.println(dp.get(n));  
    }  

    private static void solution(Long N, Long P, Long Q) {  
        if(dp.containsKey(N)) {  
            return;  
        }  

        solution(N/P, P, Q);  
        solution(N/Q, P, Q);  
        dp.put(N, dp.get(N/P) + dp.get(N/Q));  
    }  
}