[BOJ / C++] 백준 1351 무한 수열

문제

https://www.acmicpc.net/problem/1351

 

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

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

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

입력

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

출력

첫째 줄에 AN을 출력한다.

제한

• 0 ≤ N ≤ 1012
• 2 ≤ P, Q ≤ 109

---

풀이

n의 범위가 10^12까지이고, long long 타입 배열에 전체를 저장하면 문제 조건의 메모리 범위를 넘어선다.

 

아래 수열 식을 적용하되, An에서부터 구해야하는 값만 구하도록, 재귀를 사용한다.

이때 map에 {n, An}로 값을 저장하여 일반 배열처럼 사용한다. 

 

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

재귀를 사용하다보면 시간 초과가 발생한다.

⭐️ 이때, 이전 값을 꼭 저장하고, map에 계산된 값이 있으면 바로 리턴하는 로직으로 계산을 줄인다. (메모이제이션)

 

#include <cmath>
#include <iostream>
#include <map>
using namespace std;

const int MAX_N = pow(10, 12);

long long n, p, q;
map<long long, long long> a; // {n, an}

long long recursion(long long idx) {
    // 0일 때 중단
    if (idx == 0) {
        return 1;
    }
    // 이미 계산한 것이라면
    if (a.find(idx) != a.end()) {
        return a[idx];
    }
    long long idx1 = floor(idx / p);
    long long idx2 = floor(idx / q);

    // 값 저장하기
    a[idx] = recursion(idx1) + recursion(idx2);

    return a[idx];
}
int main() {

    cin >> n >> p >> q;

    a[0] = 1LL;

    cout << recursion(n);

    return 0;
}