[파이썬 | BOJ | 2331] 반복수열

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

2331번: 반복수열

문제

다음과 같이 정의된 수열이 있다.

• D[1] = A
• D[n] = D[n-1]의 각 자리의 숫자를 P번 곱한 수들의 합

예를 들어 A=57, P=2일 때, 수열 D는 {57, 74(=5^2+7^2=25+49), 65, 61, 37, 58, 89, 145, 42, 20, 4, 16, 37, …}이 된다. 그 뒤에는 앞서 나온 수들(57부터가 아니라 58부터)이 반복된다.

이와 같은 수열을 계속 구하다 보면 언젠가 이와 같은 반복수열이 된다. 이때, 반복되는 부분을 제외했을 때, 수열에 남게 되는 수들의 개수를 구하는 프로그램을 작성하시오. 위의 예에서는 {57, 74, 65, 61}의 네 개의 수가 남게 된다.

입력

첫째 줄에 A(1 ≤ A ≤ 9999), P(1 ≤ P ≤ 5)가 주어진다.

출력

첫째 줄에 반복되는 부분을 제외했을 때, 수열에 남게 되는 수들의 개수를 출력한다.

풀이

알고리즘 분류 : 그래프

• 정석 풀이

예전에 C++로 풀때 풀었던 코드이다.

계산 규칙에 따라 계속 계산하면서 DFS로 탐색해 나가면서, check배열에 지금까지 계산한 index값을 저장한다. 그러다가 이미 방문했던(계산했던) 값이 나오면 그때의 check배열에 저장된 index를 return 받아서 제출하는 방식이다.

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

int check[999999];

int Rule(int num, int P, int cnt){
    if(check[num] > 0)
        return check[num]-1;
    int sum=0;
    int temp = num;
    while(temp){
        sum += pow(temp % 10, P);
        temp /= 10;
    }
    check[num] = cnt;
    return Rule(sum, P, cnt+1);
}

int main(void){
    int num, P;
    cin>>num>>P;
    cout<<Rule(num, P, 1);
    return 0;
}

다른 풀이

새로 복습하면서 파이썬으로 풀때는 그래프 알고리즘 말고 단순한 알고리즘으로 풀어봤다.

이 알고리즘은 단순히 계산한것을 D 리스트에 넣다가, 중복이 나오면 그만두는 알고리즘으로 더 직관적이다.

import sys

A, P = map(int, sys.stdin.readline().split())
D = [A]
i = 0

while 1:
    Alist = list(str(D[i]))
    temp = 0
    for x in Alist:
        temp += int(x) ** P
    if D.count(temp) > 0:
        break
    else :
        D.append(temp)
    i += 1

Alist = list(str(D[-1]))
temp = 0

for x in Alist:
    temp += int(x) ** P
#print(temp)
try:
    print(D.index(temp))
except ValueError:
    print(0)