[백준] 2240번: 자두나무

문제

백준 - 2240번: 자두나무

분석

매 초마다, 두 개의 나무(1번, 2번) 중 하나의 나무에서 열매가 떨어지는 정보가 주어질 때, 받을 수 있는 최대 자두 개수를 구하는 문제이다. 최대 w번 움직일 수 있고, 움직이는 데 시간은 걸리지 않는다(1초보다 훨씬 짧은 시간)고 가정한다. 처음에 1번 나무에 위치한다.

움직임이 최대 w이므로 이것을 제한하기 위해 처음부터 배열크기로 제한하면 된다고 생각했다. 그래서 dp를 i는 초, j는 움직인 횟수, dp[i][j]는 (얻은 자두 개수, 마지막 위치)로 정의했다.

처음에 1번 나무에 위치하므로, 0초에 1번 나무에서 떨어진다면 가만히 있어도 자두를 1개 얻을 수 있다. 0초에 2번 나무에서 떨어진다면 움직여야 받을 수 있다.
따라서 0초에 1번 나무에서 떨어질 경우, 내가 안 움직일 때는 dp[0][0] = (1, 1)
1번 나무에 떨어지는 데 움직일 경우도 고려하면 dp[0][1] = (0, 2)가 된다.
0초에 2번 나무에서 떨어질 경우, 내가 움직일 때는 dp[0][1] = (1, 2)
2번 나무에 떨어지는 데 안 움직일 경우도 고려하면 dp[0][0] = (0, 1)가 된다.

이렇게 초기화를 하고 표를 그려가다 보면 고려해야 할 값들은 dp[i-1][j]과 dp[i-1][j-1] 이다. (움직이지 않거나, 움직이거나 둘 중 하나이기 때문)
움직이지 않았을 때, 그 전의 위치(=현재 위치)와 현재 초에 자두 나무 번호가 같다면 움직임 없이 자두 하나를 더 얻으므로 dp[i][j]에 (전 자두 개수 + 1, 전 위치)를 넣는다. 다르다면 dp[i-1][j]값을 그대로 가져간다.
움직였을 때, 그 전의 위치와 현재 초에 자두 나무 번호가 같다면
움직여서 자두 하나를 더 얻으므로 (전 자두 개수 + 1, 3 - 전 위치)를 계산한다. (3에서 1을 빼면 2, 2를 빼면 1이 되기 때문)
그렇게 했을 때 기존 dp[i][j]의 자두 개수보다 크면 갱신한다.

마지막 초까지 고려했을 때, 가장 큰 자두 개수를 구하면 된다.

코드

해결언어 : Python

import sys
input = sys.stdin.readline

t, w = map(int, input().split())
plums = []
for sec in range(t):
    num = int(input())
    plums.append(num)
dp = [[(0, 0)]*(w+1) for _ in range(t)] # i:초, j:움직임, 값:(자두, 마지막 위치)
if plums[0] == 1:
    dp[0][0] = (1, 1)
    dp[0][1] = (0, 2)
elif plums[0] == 2:
    dp[0][1] = (1, 2)
    dp[0][0] = (0, 1)

for i in range(1, t):
    for j in range(w+1):
        prev_plum, prev_pos = dp[i - 1][j]
        if plums[i] == prev_pos:
            dp[i][j] = (prev_plum+1, prev_pos)
        else:
            dp[i][j] = (prev_plum, prev_pos)
        if j:
            prev_plum, prev_pos = dp[i - 1][j - 1]
            if plums[i] == 3-prev_pos and prev_plum+1 > dp[i][j][0]:
                dp[i][j] = (prev_plum+1, 3-prev_pos)

print(max(dp[-1])[0])

끝으로..

dp 문제를 복습할 수 있던 좋은 문제였다.
dp 문제는 손으로 표를 그려보고 시작하면 좀 더 쉽게 코드를 짤 수 있는 것 같다.