백준 3866 - 풍선 수집 (Python)

문제 정보

• https://www.acmicpc.net/problem/3866
• DP

---

문제 정리

풍선이 총 N개 떨어지며, i번째 풍선은 $t_i$의 시간에 $p_i$의 위치에 떨어진다. 풍선이 땅에 닿는 순간 로봇이 같은 위치에 있어야 풍선을 잡을 수 있고, 그렇지 않으면 게임이 종료된다.
로봇은 한 번에 풍선을 최대 3개까지 들 수 있으며, 현재 들고 있는 풍선이 k개일 때 1칸 이동에 k+1초가 걸린다.
모든 풍선을 잡아 0에 위치한 집에 보관할 때까지 필요한 최소 이동거리를 구하는 문제이다.

---

풀이

딱 봐도 DP다.
일단 부분 문제는 i번째 풍선까지 처리했을 때로 볼 수 있다. 부분 문제들이 서로 이어지기 위해 필요한 정보는 로봇이 현재 들고 있는 풍선의 개수뿐이다. 이걸 상태로 잡아 점화식을 세워보자.

• DP[i][j] = i번째 풍선까지 처리했을 때, 로봇이 현재 j개의 풍선을 들고 있는 상태에서의 최소 이동거리.

DP[i]는 DP[i-1]만 알면 계산할 수 있으므로, DP 테이블을 1차원으로 줄일 수 있다.
전이 아이디어는 다음과 같다.

1. 지금 들고 있는 풍선을 유지한 채로 i번째 풍선 위치로 이동해서 잡는 경우
2. 집에 들러 들고 있는 풍선을 놓은 뒤, i번째 풍선을 잡으러 가는 경우

전이할 때 i번째 풍선이 땅에 도착하는 시각까지 행동을 처리할 수 있는지 체크해 줘야 한다.

---

후기

G2치곤 많이 할만했다. G3~G4로 봐도 될듯?
특이한 점은 하나의 보유 풍선의 개수가 3개까지라 사실상 상수 시간에 가깝고, 전체를 다 봐도 $O(N)$으로 처리할 수 있다는 점이다. 그런데 제한이 $N \leq 40$ 밖에 되지 않는다. 아마 문제에 제시되지는 않았지만 테스트케이스가 상당히 많아서 런타임이 116ms까지 나온 것 같다. 케이스를 대체 얼마나 넣어둔거지...

그리고 별개로, 한글 번역본이 조금 애매한 부분이 있어서 질문 게시판에 글을 남겨뒀다.

---

코드

# 백준 3866

import io

input = io.BufferedReader(io.FileIO(0), 1<<18).readline
INF = 10**8
MAX_X = 100


def solve(N, ball):
    prevDP = [INF] * 4
    prevDP[0] = 0
    prevT = 0
    prevX = 0

    for i in range(1, N+1):
        DP = [INF] * 4
        ballX, ballT = ball[i-1]
        canCatch = False

        # 갖고 있는 풍선을 그대로 들고 현재 풍선을 잡는 경우
        for count in range(3):
            robotT = prevT + abs(ballX - prevX)*(count+1)
            if prevDP[count] != INF and robotT <= ballT:
                canCatch = True
                DP[count+1] = min(DP[count+1], prevDP[count] + abs(ballX-prevX))

        # 갖고 있는 풍선을 집에 보관한 뒤 현재 풍선을 잡는 경우
        for count in range(1, 4):
            robotT = prevT + prevX*(count+1) + ballX
            if prevDP[count] != INF and robotT <= ballT:
                canCatch = True
                DP[1] = min(DP[1], prevDP[count] + prevX+ballX)

        # 풍선을 잡지 못한 경우
        if canCatch == False:
            return f"NG {i}"

        prevT = ballT
        prevX = ballX
        prevDP = DP

    # 최소 이동거리 도출
    result = INF
    for c in range(4):
        result = min(result, prevDP[c] + prevX)

    return f"OK {result}"


def main():
    while True:
        N = int(input())
        if N == 0:
            break
        ball = []
        for _ in range(N):
            ball.append(list(map(int, input().split())))

        print(solve(N, ball))


main()