위 문제를 간단히 요약하면 다음과 같다. 제한된 개수의 계단이 있는데, 각 계단마다 점수가 있고, 가장 마지막 계단인 도착점에 왔을 때 극대화된 점수는 몇 점인지 묻는 문제이다.
이때, 계단을 오르는 데에는 3가지 규칙이 있다.
- 계단은 한 번에 한 계단씩 또는 두 계단씩 오를 수 있다. 즉, 한 계단을 밟으면서 이어서 다음 계단이나, 다음 다음 계단으로 오를 수 있다.
- 연속된 세 개의 계단을 모두 밟아서는 안 된다. 단, 시작점은 계단에 포함되지 않는다.
- 마지막 도착 계단은 반드시 밟아야 한다.
이 규칙에 따라 극대화된 점수를 산출하면 된다.
이 문제를 시간 제한에 맞춰 풀기 위해서는 동적 계획법(Dynamic Programming) 알고리즘을 써야한다. 해당 알고리즘은 재귀(동일 함수를 계속해서 호출하는)와는 다르게, 이전 과정의 계산값을 저장한다는 점에서 중복계산을 피할 수 있다. 즉 재귀로 이 문제를 풀 시, 매 계단의 최댓값을 구할 때마다 계산이 중복되기 때문에 시간 제한을 초과하게 된다.
즉 DP란, 문제를 작은 문제들로 쪼개어 각 문제들에 대한 답을 저장해가며 문제를 푸는 알고리즘 기법이라 할 수 있다.
다음은 문제에 대한 코드이다.
import sys
input = sys.stdin.readline
# 1
N = int(input())
stairs = [int(input()) for _ in range(N)]
dp = list()
# 2
if N <= 2:
print(sum(stairs))
# 3
else:
dp.append(stairs[0])
dp.append(stairs[0]+stairs[1])
dp.append(max(stairs[0]+stairs[2], stairs[1]+stairs[2]))
for i in range(3, N):
dp.append(max(dp[i-3]+stairs[i-1]+stairs[i], dp[i-2]+stairs[i])) # 핵심 점화식
print(dp[-1])
각 주석마다 코드를 나눠서 설명하면 먼저 # 1 부분은 입력을 받는 부분이다. 이때 dp는 중간계단의 값들을 저장할 자료구조로, 리스트로 구현했다.
#2 부분은 계단이 2개 이하일 경우의 답을 출력하는 코드이다. 2개 이하이면 계단의 점수의 합이 정답이 된다. 조건이 300 이하의 자연수 개수의 계단이라고 했으니, 이러한 예외처리를 고려하는 것이 필요하겠다.
#3 부분은 본격적으로 점화식을 활용한 DP 알고리즘 구현이다. 3번째 계단까지는 하드코딩을 통해 값을 넣어줬다. 이후 4번째부터는 핵심 점화식인 max(dp[i-3]+stairs[i-1]+stairs[i], dp[i-2]+stairs[i])에 따라 값을 넣어줬다. 이 점화식에 의하면 해당 계단의 최대점수는, 이전 계단에서 바로 올라온 값과, 이전의 이전 계단에서 올라온 값을 비교하여 산출된다. 즉, 이 문제의 핵심은 모든 경우의 수를 계산하는 것이 아니라, 위에 기술된 규칙에 맞는 각 계단들의 최댓값을 중간중간 저장해가며 최종 최댓값을 구해내는 데에 있다.
재귀, 동적계획법 관련 문제들을 자꾸 마주하면서 점화식을 코드로 구현하는 연습이 계속해서 필요할 것 같다.