[BOJ C++] 1131번 - 숫자

문제

자연수 N이 주어졌을 때, N의 각 자리수를 K제곱 한 후에 그 합을 구하는 함수를 S_K(N)이라고 하자. 예를 들어, S2(65) = 62 + 52 = 61이다.
이제 다음과 같은 수열을 하나 만들어보자. N, S_K(N), S_K(S_K(N)), … 이때, A와 B와 K가 주어졌을 때, A보다 크거나 같고, B보다 작거나 같은 모든 N으로 각각 수열을 만들었을 때, 그 수열에서 가장 작은 수의 합을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 세 정수 A, B, K가 주어진다.

출력

첫째 줄에 문제의 정답을 출력한다.

제한

• 1 ≤ A ≤ B ≤ 1,000,000
• 1 ≤ K ≤ 6

문제 접근

이 문제는 DP를 이용하는 문제이다.
dp[a] = b는 S_K(a) = b이라는 뜻이다. 이렇게만 보면 가장 작은 값을 어떻게 구하나 싶은데, 사실 규칙이 있다.
각 자릿수를 K제곱해 더한 합들을 늘어놓으면 반복되는 수열의 꼴을 나타낸다.
예를 들어, 1, 5, 2 세 수가 주어졌을 때, 1~5 까지의 수들을 각 자릿수마다 제곱하여 더해나가면,

• 1: 1, 1, 1...
• 2: 2, 4, 16, 37, 58, 89, 145, 42, 20, 4...
• 3: 3, 9, 81, 65, 61, 37, 58, 89, 145, 42, 20, 4, 16, 37...
• 4: 4, 16, 37, 58, 89, 145, 42, 20, 4...
• 5: 5, 25, 29, 85, 89, 145, 42, 20, 4, 16, 37, 58, 89...

이렇게 된다.
위 수열들의 규칙은, 모두 동일한 반복수열을 가진다는 것이다. 그렇다면, 이미 한번 계산했던 것에 대해 또 S_K() 를 적용하려고 하는 경우가 무조건 있다는 말이다.
우리는 이것을 S_K()를 적용하는 반복의 종료 조건으로 달아주면 된다.
그리고, dp[] 배열에는 이전에 했던 S_K()들의 결과를 저장해 속도를 높여준다.

풀이

#include <bits/stdc++.h>
using namespace std;

int main() {
    ios_base :: sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    
	int a, b, k;
	cin >> a >> b >> k;
	
	long int answer = 0;
	vector<long int> dp(3188647, 0);
	vector<long int> visited(3188647, 0);
    for (int i = a; i <= b; i++) {
        long int myMin = i;
        int index = i;
        while(visited[index] != i) {
            visited[index] = i;
            int curNum = index;
            long int added = 0;
            if (dp[index] == 0) {
                for(int j = 1000000; j >= 1; j = j / 10) {
                    added += pow(curNum / j, k) ;
                    curNum = curNum % j;
                }
                dp[index] = added;
                if (myMin > added) myMin = added;
                index = added;
            }
            else {
                if (myMin > dp[index]) myMin = dp[index];
                index = dp[index];
            }
        }
        answer += myMin;
    }
    cout << answer;
}

 
dp는 위에서 말했듯이 S_K()들의 결과를 저장하는 배열, visited는 이전에 방문했던 수들을 기록하기 위한 배열이다.
visited 배열을 bool 값으로 만들 수도 있겠지만, 그렇게 되면 반복이 초기화 될 때마다 visited 배열 또한 초기화되어야 하기에 시간이 오래 걸리는 문제가 있었다.
그래서 초기화가 필요 없도록 현재 수를 visited 배열의 방문한 인덱스에 넣고 지금 보려는 visited배열의 인덱스에 해당하는 값이 0이 아니면 (방문했다면) 종료하는 방식이다.
최솟값을 구하는 문제이기에 최솟값 또한 dp 배열에서 찾지 않고 그 때 그 때 비교하여 저장했다.