백준: 2331 반복수열 [파이썬, 실버 4]

문제

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

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

예를 들어 A=57, P=2일 때, 수열 D는 [57, 74(=52+72=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)가 주어진다.

출력

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

예제 입력

57 2

예제 출력 

4

 

[문제 풀이]

각 자리수 0부터 9까지의 수를 거듭제곱한 합으로 수열을 만든다. 따라서 수열이 반복되는 구간이 만들어진다.

수의 중복을 확인하고, 중복 이전까지의 수열의 길이를 세어주면 된다.

초반에는 반복되는 수열의 규칙을 너무 어렵게 생각하고 아래같이 코드를 작성했었다. 아래로 작성해서 반복되는 수열을 구했는데, 인덱스만 필요해서 쉬운 문제였다.

 

틀린 코드

A = input()
P = int(input())
test=[]
answer=[]
check=0
first=-1
while True:
    num=0
    for i in range(len(A)):
        num+=int(A[i])**P
    if num in test and num!=first:
        if first==-1:
            first=num
            check=1
        answer.append(num)
        A=str(num)
    elif num==first:
        break
    else:
        test.append(num)
        A=str(num)
    
print(len(test)-len(answer)+1)

정답코드

A, P = map(int, input().split())
    
num = [A]
while True:
    tmp=0
    for i in str(num[-1]):
        tmp+=int(i)**P
    if tmp in num:
        break
    num.append(tmp)
    
print(num.index(tmp))