[DP] 백준 2240번 “자두나무” Python 풀이

문제 링크

• https://www.acmicpc.net/problem/2240

풀이

• 예제 입력1에 대한 DP 테이블은 아래와 같이 구성할 수 있습니다.

  

  • 행 : 지난 초 (i)
  • 열 : 사람이 움직인 횟수 (j)
  • dp[i][j] : i초까지 j번 움직였을 때, 얻을 수 있는 최대 자두 수

• DP 테이블은, i초에 자두를 먹을 수 있는 조건과 먹을 수 없는 조건을 통해 채울 수 있습니다.

  • 자두를 먹을 수 있는 조건
    • i초에 자두가 2번 나무에서 떨어지고, 사람의 이동 횟수가 홀수라면(j % 2 == 1), 현재 2번 나무에 위치해 있다는 뜻이므로, 자두를 받아먹을 수 있음
      • dp[i][j] = max( i-1초 때 j 열의 자두 수, i-1초 때 j-1 열의 자두 수 ) + 1
      • 이전 위치로부터 움직여서 받아 먹을 것인지, 현재 위치에서 받아 먹을 것인지를 비교
    • i초에 자두가 1번 나무에서 떨어지고, 사람의 이동 횟수가 짝수라면(j % 2 == 0), 현재 1번 나무에 위치해 있다는 뜻이므로, 자두를 받아 먹을 수 있음
      • dp[i][j] = max( i-1초 때 j 열의 자두 수, i-1초 때 j-1 열의 자두 수 ) + 1
      • 이전 위치로부터 움직여서 받아 먹을 것인지, 현재 위치에서 받아 먹을 것인지를 비교
  • 자두를 먹을 수 없는 조건
    • i초에 떨어지는 자두와 현재 사람의 위치가 다를 때
      • dp[i][j] = max( i-1초 때 j 열의 자두 수, i-1초 때 j-1 열의 자두 수 )
      • 먹지는 못하므로, 움직여서 못 먹는 경우와 움직이지 않아서 못 먹는 경우를 비교

• 결과적으로, DP 테이블의 마지막 행의 최댓값을 출력하면 정답이 됩니다.

코드

import sys

# 자두가 떨어지는 시간 T와 이동할 수 있는 횟수 W
t, w = map(int, sys.stdin.readline().rstrip().split(" "))

# 자두가 떨어지는 나무의 번호 저장
arr = [0]
for _ in range(t):
    arr.append(int(sys.stdin.readline()))

# 2차원 DP 테이블 초기화
dp = [[0] * (w+1) for _ in range(t+1)]

for i in range(t+1):

    # 1번 나무에서 한 번도 움직이지 않는 경우

    # 1번 나무에서 자두가 떨어진다면
    if (arr[i] == 1):
        dp[i][0] = dp[i-1][0] + 1

    # 2번 나무에서 자두가 떨어진다면
    else:
        dp[i][0] = dp[i-1][0]

    # 1번 이상 움직이는 경우에 대해 탐색
    for j in range(1, w+1):

        # i초에 2번 나무에서 자두가 떨어지고, 현재 2번 나무에 위치해있다면
        if (arr[i] == 2 and j % 2 == 1):
            # 이전 위치로부터 움직여서 받아 먹을 것인지, 현재 위치에서 받아 먹을 것인지를 비교
            dp[i][j] = max(dp[i-1][j-1], dp[i-1][j]) + 1

        # i초에 1번 나무에서 자두가 떨어지고, 현재 1번 나무에 위치해있다면
        elif (arr[i] == 1 and j % 2 == 0):
            # 이전 위치로부터 움직여서 받아 먹을 것인지, 현재 위치에서 받아 먹을 것인지를 비교
            dp[i][j] = max(dp[i-1][j-1], dp[i-1][j]) + 1

        # i초에 자두가 떨어지는 나무와 현재 나무의 위치가 다르다면
        else:
            # 움직여서 못 먹는 경우와 움직이지 않아서 못 먹는 경우를 비교
            dp[i][j] = max(dp[i-1][j-1], dp[i-1][j])

# DP 테이블의 마지막 행의 최댓값을 출력
print(max(dp[t]))

참고

• https://deok2kim.tistory.com/101
• https://codedrive.tistory.com/438