[Python] 백준 #2331 반복수열

www.acmicpc.net/problem/2331

2331번: 반복수열

처음 코드에서 런타임 오류가 뜨길래 코드 자체의 오류라기보단, 범위 설정이 잘못되었다고 생각하여 문제 조건을 다시 읽어보았다

 

1<=A<=9999
1<=P<=5 이므로 방문을 확인하는 visited 배열의 범위를 (9^5)*4로 설정을 해주니 통과!

 

# 1<=A<=9999
# 1<=P<=5

from collections import deque
import sys

que = deque()
k=(9**5)*4
visited = [0] * (k+1)

A, P = map(int, sys.stdin.readline().split())

i = A
que.append(i)
visited[i] = 1

check = 0
cnt = 0

while True:
    data = 0

    for j in str(i):
        data += int(j) ** P

    i = data

    if visited[i] == 0:
        que.append(i)
        visited[i] = 1

    else:  # 이미 한번 방문했었음
        visited[i] = 2
        break

while que:
    p = que.popleft()
    if visited[p] == 2:
        break
    elif visited[p] == 1:
        cnt += 1

print(cnt)