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

출력

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

---

풀이 과정

1. set에 반복되는 요소가 나올때까지 수열을 넣는다
   • 반복되는 요소가 나온다면 해당 요소의 값을 따로 저장만 하고 넣지는 않음.
2. 반복되는 요소에서부터 set에서 나오지 않을때까지 수열을 구하고 해당 수열을 set에서 빼준다.
3. set에 남아있는 수열이 즉, 반복되는 부분이 제외된 수열이다.

---

소스 코드

import sys
input = sys.stdin.readline

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

def get_next(n, p):
    nvalue = 0
    while n > 0:
        nvalue += (n % 10) ** p
        n = n // 10
    return nvalue

visited = set()

current = A
dup_start = 0
# 1. 지나간 숫자중에 하나가 다시 나타나는 경우 중복의 시작
while True:
    if current not in visited:
        visited.add(current)
    else:
        dup_start = current
        break
    current = get_next(current, P)

# 2. 중복 시작점에서 중복 끝날때까지 요소들 빼줌
while dup_start in visited:
    visited.remove(dup_start)
    dup_start = get_next(dup_start, P)

# 3. 남아있는 요소의 개수 출력
print(len(visited))