[Python] 포도주 시식

백준 2156번

문제 설명

효주는 포도주 시식회에 갔다. 그 곳에 갔더니, 테이블 위에 다양한 포도주가 들어있는 포도주 잔이 일렬로 놓여 있었다. 효주는 포도주 시식을 하려고 하는데, 여기에는 다음과 같은 두 가지 규칙이 있다.

1. 포도주 잔을 선택하면 그 잔에 들어있는 포도주는 모두 마셔야 하고, 마신 후에는 원래 위치에 다시 놓아야 한다.
2. 연속으로 놓여 있는 3잔을 모두 마실 수는 없다.

효주는 될 수 있는 대로 많은 양의 포도주를 맛보기 위해서 어떤 포도주 잔을 선택해야 할지 고민하고 있다. 1부터 n까지의 번호가 붙어 있는 n개의 포도주 잔이 순서대로 테이블 위에 놓여 있고, 각 포도주 잔에 들어있는 포도주의 양이 주어졌을 때, 효주를 도와 가장 많은 양의 포도주를 마실 수 있도록 하는 프로그램을 작성하시오.

예를 들어 6개의 포도주 잔이 있고, 각각의 잔에 순서대로 6, 10, 13, 9, 8, 1 만큼의 포도주가 들어 있을 때, 첫 번째, 두 번째, 네 번째, 다섯 번째 포도주 잔을 선택하면 총 포도주 양이 33으로 최대로 마실 수 있다.

풀이

dp 문제를 풀 때는 항상 점화식을 위해 조건을 먼저 생각해야 하는데
연속으로 세 잔 마실 수 없다고 했으니 마지막 잔 기준으로 생각하면 대충 조건이 나온다.

1. 마지막 잔을 마시고 이전 잔도 마셨을 경우
2. 마지막 잔을 마시고 이전 잔을 마시지 않았을 경우
3. 마지막 잔을 마시지 않았을 경우

다음 조건들을 수식으로 나타내면 다음과 같다.

1. wine[i] + dp[i - 1] + dp[i - 3]
2. wine[i] + dp[i - 2]
3. dp[i - 1]

점화식은 위 세 가지 조건의 max값으로 만들면 된다.

dp[i] = max(wine[i] + dp[i - 2], wine[i] + wine[i - 1] + dp[i - 3], dp[i - 1])

오답

n = int(input())
dp = [0] * n
wine = [int(input()) for _ in range(n)]

dp[0] = wine[0]
dp[1] = wine[0] + wine[1]
dp[2] = max(wine[2] + wine[0], wine[2] + wine[1], dp[1])

for i in range(3, n):
    # 마지막 잔을 마시지 않을 경우도 넣어야 한다.
    dp[i] = max(wine[i] + dp[i - 2], wine[i] + wine[i - 1] + dp[i - 3], dp[i - 1])

print(dp[n - 1])

처음에 이렇게 제출했다가 런타임 에러(IndexError)로 털렸다.
원인은 dp[1], dp[2] 정의 부분이었는데
생각해보니 주어지는 n이 1일수도 있는데 (즉 dp의 길이가 1일 수도 있는데) 아무 조건 없이 와인잔이 2개, 3개일 때를 가정해서 틀렸던 것이었다.

제출 코드

n = int(input())
dp = [0] * n
wine = [int(input()) for _ in range(n)]

dp[0] = wine[0]

if n > 1:
    dp[1] = wine[0] + wine[1]

if n > 2:
    dp[2] = max(wine[2] + wine[0], wine[2] + wine[1], dp[1])

for i in range(3, n):
    # 마지막 잔을 마시지 않을 경우도 넣어야 한다.
    dp[i] = max(wine[i] + dp[i - 2], wine[i] + wine[i - 1] + dp[i - 3], dp[i - 1])

print(dp[n - 1])

그래서 if문을 추가해서 해결.