[백준] 15678 - 연세워터파크 (Python)

문제의 조건을 보면, 아무 지점에서 시작해서 나오고 싶을 때 나오는 것이 가능하다. 그렇다면, 나오는 지점에 초점을 두고 동적 계획법을 적용해 보자.

시작점과 도착점을 반대로 바꿔도 점수에는 변화가 없으므로 항상 왼쪽에서 시작해 오른쪽에서 탈출한다고 가정해도 문제가 없다. 따라서 오른쪽 시작점은 고려하지 않는다.

$i$보다 왼쪽에 있는 어떠한 시작점에서 출발해 도착점을 $i$로 할 때의 최댓값을 $DP[i]$라 하자.

최대로 건너뛸 수 있는 값을 $D$, 징검다리 배열을 $K$라고 하자. 점화식을 의사코드로 표현하면 아래와 같다.

$DP[i] = max(DP[i-D], ..., DP[i-1]) + K[i]$

단순히 반복문으로 구현하면 $O(N^2)$으로 비효율적이다.

이 문제의 핵심은 구간 내의 최댓값을 빠르게 찾는 것이므로 세그먼트 트리를 사용하면 $O(NlogN)$으로 개선이 가능하다.

import math

# 구간 내 최댓값 갱신
def make_tree(left, right, node, idx, value):
    seg_tree[node] = max(seg_tree[node], value) 
    if left == right:
        return
    mid = (left+right) // 2 
    if idx <= mid:
        make_tree(left, mid, node*2, idx, value)
    else:
        make_tree(mid+1, right, node*2+1, idx, value)

# [start, end] 구간 내의 최댓값 반환
def find_max(left, right, node, start, end):
    if start > right or end < left:
        return -math.inf
    if left == right:
        return seg_tree[node]
    if start <= left and right <= end:
        return seg_tree[node]
    mid = (left+right) // 2 
    return max(find_max(left, mid, node*2, start, end), find_max(mid+1, right, node*2+1, start, end))
    
n, d = map(int, input().split())
# 인덱스를 1부터 사용하기 위함
dp = [0]
# 시작한 지점에서 바로 나가는 경우가 초기값
dp.extend(map(int, input().split()))
seg_tree = [-math.inf]*(n*4)
for i in range(1, n+1):
	# 최댓값 갱신
    dp[i] = max(dp[i], dp[i]+find_max(1, n, 1, i-d, i-1))
    # 만들어진 최댓값을 세그먼트 트리에 삽입
    make_tree(1, n, 1, i, dp[i])
print(max(dp[1:]))

덱을 사용하면 더 빠른 $O(N)$ 구현이 가능하다.
덱에는 값이 $(i, DP[i])$형태로 들어가며 항상 $DP[i]$값에 대해 내림차순을 유지해야 한다.

from collections import deque

n, d = map(int, input().split())
dp = [0]
dp.extend(map(int, input().split()))
deq = deque()
for i in range(1, n+1):
    while deq:
    	# i-d 보다 왼쪽의 값은 삭제
        if deq[0][0] < i-d:
            deq.popleft() 
        else:
        	# 최댓값은 첫번째 값
            dp[i] = max(dp[i], dp[i]+deq[0][1])
            break
    while deq:
    	# 뒤에서부터 현재 값보다 작은 값은 모두 삭제
        if deq[-1][1] < dp[i]:
            deq.pop()
        else:
            break
    deq.append((i, dp[i]))
print(max(dp[1:]))