백준 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)가 주어진다.

출력

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

---

정답

from collections import defaultdict

A, P = map(int, input().split())

d = []
d.append(A)

count = defaultdict(int)
count[A] += 1

while True:
    n = d[-1]
    new = sum([int(i) ** P for i in str(n)])
    d.append(new)
    count[new] += 1
    if count[new] == 2:
        break

print(d.index(new))

 

---

풀이

from collections import defaultdict

A, P = map(int, input().split())

d = []
d.append(A)

count = defaultdict(int)
count[A] += 1

while True:
    n = d[-1]
    new = sum([int(i) ** P for i in str(n)])
    d.append(new)
    count[new] += 1
    if count[new] == 2:
        break

print(d.index(new))

 이 문제는 정의된 수열에서 반복되는 부분을 찾아서 제거하면 되는 문제입니다.

 해당 수열은 각 자릿수의 합으로 이루어진 수열인데 여기서 반복되는 규칙을 먼저 찾아야 합니다. 같은 숫자가 두 번 이상 나오면 무조건 반복이 되는 것이므로, 숫자별로 나온 횟수를 세어가면서 두 번째 나온 순간 반복문을 멈춘 다음 첫 번째 인덱스를 반환하면 됩니다.