알고스팟 원주율 외우기 (PI) 문제 파이썬으로 풀기

숫자를 입력받아, 난이도를 계산하는 문제이다.

 

전형적인 동적계획법 문제처럼 보이지만(큰 문제를 작게 자르고 값을 계산하여 저장하고 저장한 값을 이용해 최종 값을 구하면 된다), 이 문제의 특이한 점은 숫자들을 3, 4, 5자리로만 끊어야 하고, 입력이 매우 크다는 것이다.

두번째 점은 특히 파이썬으로 풀때 치명적으로 다가오는데, 일반적인 동적계획법 문제 풀듯이

전체 값을 저장할 리스트/배열을 만들고

리스트/배열의 특정 위치의 값을 가져와서

그 값의 최소/최댓값을 구하고 그것을 return 하는 방식으로는 주어진 시간 (1000ms) 안에 풀 수 없다는 것이다.

import sys
sys.setrecursionlimit(10**8)
 
 
def difficulty(start, end):
    part = number[start:end+1]
    if part == [part[0] for k in range(len(part))]:
        return 1
    onepro = True
    for i in range(len(part) - 1):
        if part[i + 1] - part[i] != part[1] - part[0]:
            onepro = False
            break
    if onepro and (abs(part[1] - part[0]) == 1):
            return 2
    alter = True
    for j in range(len(part)):
        if part[j] != part[j % 2]:
            alter = False
            break
    if alter:
        return 4
    if onepro:
        return 5
    return 10
 
 
def pi(pos):
    if pos == len(number):
        return 0
    ret = memo[pos]  # 특정 위치의 값을 가져와서
    if ret != -1:  # 이미 구한 값인지 알아보고
        return ret
    ret = maxnum  
    for o in range(3, 6):
        if pos + o <= len(number):
            ret = min(ret, pi(pos + o) + difficulty(pos, pos + o - 1))  # 구하지 않은 값이면 최솟값을 구하고 
    return ret  # 구한 값을 return
 
 
memo = [-1] * 10001  # 전체 값을 저장할 리스트
maxnum = 999999999
case = int(input())
for k in range(case):
    number = list(map(int, list(input().strip())))
    print(pi(0))

즉 이렇게 동적계획법 문제를 푸는 정석답게 코드를 짜면,

시간초과가 난다는 것이다.

 

그래서 c++로 우선 문제를 풀고, 파이썬으로 푼 다른 사람의 코드를 보고 정답이 뜨게 되는 코드를 작성했다.

# pi
def score_3(a, b, c):  # 3자리씩 끊어 읽기
    if a == b == c:  # 모든 숫자가 같음
        return 1
    if a-b == b-c == 1 or a-b == b-c == -1:  # 단조증가/감소
        return 2
    if a == c != b:  # 숫자가 번갈아 출현
        return 4
    if a-b == b-c:  # 등차수열
        return 5
    return 10
 
 
def score_4(a, b, c, d):  # 4자리씩 끊어 읽기
    if a == b == c == d:
        return 1
    if a-b == b-c == c-d == 1 or a-b == b-c == c-d == -1:
        return 2
    if a == c and b == d and a != b:
        return 4
    if a-b == b-c == c-d:
        return 5
    return 10
 
 
def score_5(a, b, c, d, e):  # 5자리씩 끊어 읽기
    if a == b == c == d == e:
        return 1
    if a-b == b-c == c-d == d-e == 1 or a-b == b-c == c-d == d-e == -1:
        return 2
    if a == c == e and b == d and a != b:
        return 4
    if a-b == b-c == c-d == d-e:
        return 5
    return 10
 
 
def pi(arr):
    N = len(arr)  # 리스트의 길이
    cache = [None] * (N+1)
    cache[3] = score_3(arr[0], arr[1], arr[2])  # 0~2번까지 3개 읽음
    cache[4] = score_4(arr[0], arr[1], arr[2], arr[3])  # 0~3번까지 4개 읽음
    cache[5] = score_5(arr[0], arr[1], arr[2], arr[3], arr[4])  # 0~4번까지 5개 읽음
    for i in range(6, N+1):
        cand = []
        if cache[i-3] is not None:
            cand.append(cache[i-3] + score_3(arr[i-3], arr[i-2], arr[i-1]))
        if cache[i-4] is not None:
            cand.append(cache[i-4] + score_4(arr[i-4], arr[i-3], arr[i-2], arr[i-1]))
        if cache[i-5] is not None:
            cand.append(cache[i-5] + score_5(arr[i-5], arr[i-4], arr[i-3], arr[i-2], arr[i-1]))
        cache[i] = min(cand)  # 최소값을 cache에 넣는다
    return cache[-1]
 
 
case = int(input())
for k in range(case):
    number = list(map(int, list(input().strip())))  # 입력을 숫자 리스트로 만든다
    print(pi(number))

코드를 참고한 출처: https://algospot.com/judge/submission/detail/622178

 

이 코드는 3자리, 4자리, 5자리씩 끊어 읽는 것을 함수로 각각 구현하고, 한 함수 (코드에서는 pi 함수) 내에서 값을 저장할 리스트 (cache) 를 만들고, 이를 이용해 최솟값을 구하는 방식으로 이루어졌다. 작은 입력 뿐만 아니라 긴 입력도 빠르게 결과를 나오게 해주는 아주 잘 짜여진 코드라고 생각한다.

 

느낀점:

1. 문제를 풀며 파이썬에 재귀 횟수 제한이 있다는 것을 알게 되었다. 그래서 sys.setrecursionlimit 함수를 써서 그것을 풀어주기도 했지만, 별 도움은 안되었던 것 같다.

2. 정답이 뜬 코드도 cache의 0,1,2번째는 값이 none에서 바뀌지 않고 그대로 있어서, 불필요하다고 생각해 없애고 맨 처음부터 숫자가 들어가게 하려고 했는데, 여러 방면으로 시도해보았지만 오류나고 답이 이상하게 나와서 고치는 것을 포기하게 되었다. 내가 좀 더 실력이 있었다면 바꿀 수 있었을 텐데 아쉽다.