https://www.acmicpc.net/problem/5681
| 시간 제한 | 메모리 제한 |
| 1초 | 128MB |
문제
KDK방송국은 새로운 게임 쇼를 하나 만들었다. 참가자는 선택을 몇 개 하게 되고, 이 선택에 따라서 상품을 얻게 된다.
먼저, 공이 삼각형 모양으로 쌓여져 있고, 각 공에는 정수 값이 하나씩 쓰여 있다. 아래 그림은 한 예이다.
참가자는 공을 고를 수 있고, 고른 공에 쓰여 있는 숫자의 합이 점수가 된다. 공을 고르면, 그 공은 삼각형에서 제거된다. 점수가 높을수록 좋은 상품을 받게 된다. 하지만, 참가자는 그 공의 위에 있는 공을 고른 경우에만 그 공을 고를 수 있다. 또, 참가자는 공을 고를 것인지, 게임을 중단할 것인지 선택할 수 있다. 만약, 공을 하나도 고르지 않은 경우에 점수는 0이 된다.
프로그램 PD 김동규는 참가자들이 얻을 수 있는 점수의 최댓값을 구해보려고 한다. 점수의 최댓값은 몇 점일까?
입력
입력은 여러 개의 테스트 케이스로 이루어져 있다. 각 테스트 케이스의 첫째 줄에는 공이 총 몇 행으로 쌓여져 있는지가 주어진다. 이 값을 N이라고 한다. (1 ≤ N ≤ 1000) 다음 N개의 줄의 i번째 줄에는 총 i개의 정수 Bij가 주어진다. (-105 ≤ Bij ≤ 105, 1 ≤ j ≤ i ≤ N) Bij는 i행 j열에 있는 공에 쓰여 있는 정수이다. (첫 행은 가장 윗 행, 각 행의 첫 번째 공은 가장 왼쪽 공)
입력의 마지막 줄에는 0이 하나 주어진다.
출력
각 테스트 케이스에 대해서, 참가자가 얻을 수 있는 점수의 최댓값을 출력한다.
풀이
이거 남미 ICPC 문제 번역한 건데 문제가 조금 이해하기 어렵습니다. [그 공의 위에 있는 공을 고른 경우에만 그 공을 고를 수 있다.] 부분이 무슨 말인가 싶은데 배치를 보면 왼쪽, 오른쪽 변에 위치한 공을 제외한 모든 공은 위에 두 개의 공이 있습니다. 왼쪽, 오른쪽 변에 있는 공들은 한쪽이 없으니 한 개, 맨 위 꼭대기에 있는 공은 위에 공이 없습니다. 가장 꼭대기에 있는 공은 위에 공이 없으니 그냥 선택할 수 있습니다. 왼쪽, 오른쪽 변에 있는 공들은 바로 위에 있는 하나의 공이 선택되었다면 선택할 수 있습니다. 나머지는 위에 있는 양쪽 공이 모두 선택되어야 선택할 수 있습니다.
예제 입력 1의 결과는 위 그림과 같습니다. 파란 공들을 쭉 선택하면 총 합이 7점으로 최대입니다. 가장 큰 점수인 9점을 고르려면 그 위에 있는 -8점과 2점을 골라야 합니다. 2점을 고르려면 그 위에 -5점과 3점을 골라야 합니다. 그러려면 맨 위의 3점도 골라야합니다. 그러고 나면 -8점 아래의 3점은 그냥 고를 수 있습니다. 오른쪽 맨 밑의 7점을 먹으려면 그 위에 있는 -8점을 골라야하는데 그러면 합이 -1로 손해이기 때문에 고르지 않습니다. 필요한 아이디어를 찾기 위해 그림을 좀 크게 그려서 살펴봤습니다.
빨간 공을 선택하려면, 초록 공들이 모두 선택되어 있어야 합니다. 뭔가 누적합으로 쉽게 해결될 것처럼 생겼지만, 오른쪽 처럼 선택할 경우, 단순한 누적합으로 해결하기 어려워집니다. 특정 공을 선택하려면 왼쪽과 오른쪽 위에 있는 공이 선택되어 있어야 한다는 조건이 좀 까다롭습니다. 이 중 하나를 먼저 선택 된 상태로 바꿔놓고 오른쪽 공을 선택하거나 하지 않는 쪽으로 DP를 실행해보겠습니다.
빨간 공을 고르는 상황을 생각해봅시다. 빨간 공의 위치는 (i, j)로 표현하면 (4, 1)입니다. 빨간 공을 선택하려면 파란 공들을 모두 선택한 상태여야 하고, 왼쪽의 초록 공들은 선택해도, 안해도 됩니다. 현재 확인하고 있는 j열이 1번이므로 0번 열에서 누적합을 생성해야 합니다. j = 0열에서 i = 3~7행까지의 누적합 중, 가장 큰 값을 불러와 j = 1열의 i = 4까지의 누적합에 더합니다. 이걸 dp[4][1]에 저장합니다.
여기서 가장 큰 값을 불러오기 위해 우선순위 큐를 사용합니다. j열을 확인할 때, j-1열의 모든 누적합과 행번호를 우선순위 큐에 저장하고, j열의 i행을 확인할 때는 i - 2행까지의 값을 모두 큐에서 pop한 다음 먼저 나오는 가장 큰 값을 더하는 식으로 관리합니다.
# hq = {(-prefix_sum, i) | prefix_sum = 이전 열 i행까지의 누적합}
while hq[0][1] < i - 1:
heappop(hq)
dp[i][j] = prifix_sum[i][j] - hq[0][0]
이렇게 하나의 열이 완성되면, 우선순위 큐를 모두 비운 다음 그 열의 dp값을 다시 우선순위 큐에 push합니다. 모든 공에 대해 이 작업을 수행하는 동안, 가장 큰 dp값이 최적해가 됩니다.
정답 코드
import sys
from heapq import heappush, heappop
input = sys.stdin.readline
def solution():
while True:
n = int(input())
if n == 0: break
prefix = []
for i in range(n):
prefix.append(list(map(int, input().split())))
for j in range(i):
prefix[i][j] += prefix[i - 1][j]
dp = [[0] * n for _ in range(n)]
hq = []
for i in range(n):
dp[i][0] = prefix[i][0]
heappush(hq, (-prefix[i][0], i))
result = -hq[0][0]
if result < 0: result = 0
for j in range(1, n):
for i in range(j, n):
while hq[0][1] < i - 1:
heappop(hq)
dp[i][j] = prefix[i][j] - hq[0][0]
if result < dp[i][j]:
result = dp[i][j]
hq.clear()
for i in range(j, n):
heappush(hq, (-dp[i][j], i))
print(result)
return
solution()
