백준 31838 - 아이템 2 (Python)

문제 정보

• https://www.acmicpc.net/problem/31838
• DP, 누적 합

---

문제 정리

무한히 긴 수직선 위에 N개의 아이템이 구간 [1, N]에 떨어져 있으며, 각 아이템의 가치는 $A_i$의 값을 가진다.
또한, 아이템을 주울 수 있는 집게가 있는데, 길이가 K인 집게를 위치 x에서 사용할 경우 구간 [x, x+K-1] 범위의 아이템을 전부 줍는다.
이 집게를 원하는 만큼 사용하여 주운 아이템 가치의 합이 최대가 되도록 해야 하며, 그때의 가치 최대값을 구하는 문제이다.

---

풀이

• 부분 문제로 나눌 수 있어 보인다.
  현재 좌표 x에서 할 수 있는 경우의 수는 <집게를 사용하는 것>과 <집게를 사용하지 않는 것> 2가지이다.

• 각 케이스별로 계산해보자.

  • 먼저 집게를 사용할 경우, [x-K+1, x] 구간의 아이템을 줍는 것으로 볼 수 있다.

    • <구간 [1, x-K]까지의 최대값> + <구간 [x-K+1, x]의 합>

  • 집게를 사용하지 않은 경우, 최대값은 다음과 같다.

    • <구간 [1, x-1]까지의 최대값>

• 이것만 고려해서는 예제 1을 제대로 처리할 수 없다.
  예제 1의 3, -1, 5를 집게를 사용해 집고, (-1), (5), 4를 집게를 사용해 집었을 때 나오는 가치 11이 최대값이기 때문이다.
  즉, 연속으로 집게를 사용할 경우 처리를 해줘야 한다.

• 추가해줘야 하는 내용은 다음과 같다.

  • 만약, x-1칸까지 아이템을 주워가도록 집게를 이미 사용했고, x칸까지 아이템을 줍도록 집게를 사용한 경우엔 다음과 같이 계산할 수 있다.
    • <구간 [1, x-1]까지의 최대값> + <위치 x의 아이템 가치>

• 이처럼 크게 "집게를 사용하는 경우", "집게를 사용하지 않는 경우"로 나누고, 연속해서 집는 케이스까지 생각하여 DP 테이블을 계산해주면 된다.
  DP 테이블에 저장할 값은 다음과 같다.

  • DP[x][i] = (현재 구간에서 아이템을 주웠는지 여부가 i일 때, 구간 [1, x]에서 획득할 수 있는 가치의 최대값)

• 1번 케이스에서 구간의 합을 구해야 하는데, 값이 수정되는 온라인 쿼리가 아니기 때문에 누적합 배열을 미리 전처리 해두면 $O(1)$에 구간합 쿼리를 해결할 수 있다.

• 구간 [1, N] 밖에서도 집게를 사용할 수 있으므로, 누적 합 배열을 이용해 구간 합을 구할 때 범위에 주의해줘야 한다.

---

후기

엣지케이스가 없는건지 예제가 친절한 건지는 모르겠는데 바로 맞았다.

---

코드

# 백준 31838

'''
DP[i][isCatch] = i번째 아이템을 주웠는지 여부가 isCatch일 때, [0, i] 구간에서의 가치 최대값.
'''

import io

input = io.BufferedReader(io.FileIO(0), 1<<18).readline


def solve(N, K, A):
    # 누적 합 배열 전처리
    accSum = [0 for _ in range(N)]
    accSum[0] = A[0]
    for i in range(1, N):
        accSum[i] = accSum[i-1] + A[i]
    
    # DP 계산
    DP = [[0, 0] for _ in range(N+K)]
    DP[0][0] = 0
    DP[0][1] = A[0]

    for i in range(1, N+K):
        # 이전 칸에서 집고 이번 칸에서 또 집는 경우
        DP[i][0] = max(DP[i-1])
        if i < N:
            DP[i][1] = DP[i-1][1] + A[i]

        # K칸의 아이템을 집는 경우
        if K <= i < N:
            DP[i][1] = max(DP[i][1], max(DP[i-K]) + accSum[i] - accSum[i-K])

        # N칸을 넘어서 주워야 하는 경우
        if i > N:
            DP[i][1] = max(DP[i-K]) + accSum[N-1] - accSum[i-K]

    return max([max(i) for i in DP])


def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    print(solve(N, K, A))


main()