[백준] 2331 - 반복수열 (python 파이썬)

백준 2331 문제 바로가기

풀이 코드

n, m = map(int, input().split())
check = [0] * 236197  # 9 ** 5 + 9 ** 5 + 9 ** 5 + 9 ** 5 최대 인덱스의 값이기 때문
iteration = 1


def cal(a, b):
    result = 0
    for i in str(a):
        result += pow(int(i), b)
    return result


def dfs(n, m, iteration, check):
    if check[n] != 0:
        return check[n] - 1

    check[n] = iteration
    iteration += 1
    result = cal(n, m)
    return dfs(result, m, iteration, check)


print(dfs(n, m, iteration, check))

dfs 알고리즘을 사용해서 풀었다. 우선, 각 자릿수를 p만큼 제곱하는 함수가 필요하다고 생각되어 문자열로 변환해서 for문을 통해 각 자릿수를 p만큼 다시 정수형으로 변환해 계산한 후 리턴하도록 했다.

예를 들면, 57을 문자열로 변환해 for문으로 i = 5인 상태에서 pow 함수를 사용하여 다시 i를 정수형으로 변환해 p만큼 제곱해주도록 하였다.

check배열을 통해 기존에 탐색이 되었던 값인지 확인하였는데 최대 크기는 9999의 각 자릿수들을 5제곱한 값이다.
시작점은 57이기 때문에 check[57]을 0에서 1로 바꿔주는데 처음에는 dfs로 풀려고 해서 기계적으로 탐색한 부분을 전부 1로 바꿀려고 했다
그런데 그렇게 풀면 몇번 반복되었는지에 대한 값을 추출할 수 없었다.

그래서 iteration 변수에 1을 초기값으로 시작하여 +1로 늘려가고 check인덱스 값에 대입해주며 반복하다가 같은 값이 나왔을 때, return하도록 했다.

회고

이 문제도 바로바로 풀이가 튀어나오지는 않았지만 그래도 알고리즘을 이해하다보니 풀이를 전부 지우고 천천히 풀어나가는데 막힘이 없었다!
dfs 문제도 얼른 정복하고 싶다...