[Python/파이썬] 백준 14452번 Cow Dance Show

https://www.acmicpc.net/problem/14452

코드

from heapq import heappop, heappush
import sys
input = sys.stdin.readline

n, t_max = map(int, input().split())
durations = [int(input()) for _ in range(n)]


def calculate_total_duration(k):
    stage = []
    # 시작
    for i in range(k):
        heappush(stage, durations[i])

    # 무대 진행
    time = 0
    d_idx = k
    while time <= t_max:
        # 공연 다 한 소 내려와
        while stage and stage[0] == time:
            heappop(stage)
        # 자리 남았으면 소 올라가
        while len(stage) < k and d_idx < n:
            heappush(stage, time+durations[d_idx])
            d_idx += 1
        # 소들 공연 끝났으면 그 때의 시간 리턴
        if len(stage) == 0:
            return time
        time += 1
    # t_max 내로 공연을 못 마친 경우
    return t_max+1


left, right = 1, n
while left <= right:
    k = (left+right)//2
    needed_time = calculate_total_duration(k)
    if needed_time <= t_max:
        right = k-1
    else:
        left = k+1
print(right+1)

 

보통 이렇게 정해진 범위 안에서 가장 작은 값 또는 가장 큰 값을 구하는 문제는 이분탐색(더 정확하게 말하자면 파라메트릭 서치)를 사용하면 된다.

 

이 문제도 1 ~ N 사이의 수 중에 T_max 시간 이내로 무대를 마칠 수 있는 최소의 K를 구하는 문제이기 때문에 파라메트릭 서치로 풀 수 있다.

 

이분탐색

1. K가 될 수 있는 범위는 1~N이므로 left와 right의 초기 값을 각각 1과 N으로 해둔다.

2. left와 right의 중간 값인 K를 무대의 크기로 했을 때 N마리의 소가 T_max 시간 내로 무대를 마칠 수 있는지 확인한다. 이것은 `calculate_total_duration()`를 사용하여 구한 소요 시간이 T_max 이하인지 확인하면 된다.

  2-1) 만약 T_max 시간 내로 무대를 마칠 수 있다면 right = k-1으로 조정하여 가능한 더 작은 값이 있는지 확인한다.

  2-2) 만약 T_max 시간 내로 무대를 마칠 수 없다면 left = k+1로 조정하여 무대의 크기를 키운다.

3. 이분 탐색이 끝나면(left > right가 되면), right+1이 우리가 찾는 최소 K 값이 된다. 왜냐하면, 마지막 반복에서, right는 가능한 값에서 -1된 상태로 끝나기 때문이다. 따라서 right+1은 모든 소가 T_max 시간 내에 무대를 마칠 수 있는 최소 무대 크기가 된다.

 

calculate_total_duration()

1. 처음에 올라갈 K마리의 소를 stage라는 heapq에 넣는다. stage에는 해당 소가 무대를 마치는 시간을 넣을 것인데, 초기 상태에서는 0초에서 시작하는 것이므로 durations에 있는 수를 그대로 넣으면 된다.

2. 최대 T_max초까지 다음 과정을 반복하며 모든 소가 무대를 마치고 내려갈 수 있는 시간을 찾는다.

  2-1) stage는 가장 작은 값이 앞에 오는 최소힙이다. 그렇기 때문에 stage[0]에는 무대에 내려가는 시간이 가장 빠른 소가 오게 된다. 가장 앞의 소부터 확인하며 내려갈 시간이 된 소들을 heappop한다.

  2-2) stage의 길이가 k 미만이라면 자리가 남았으므로 소들이 더 올라갈 수 있다. 길이가 k가 될 때까지 소들을 올려보낸다. 이때 heappush되는 값은 (현재 시간+공연 시간)으로 넣는다.

  2-3) 만약 2-1과 2-2 과정을 거쳤는데 stage가 비었다면 모든 소가 공연을 끝낸 것이다. 이때의 time을 리턴해주면 된다.

  2-4) 시간을 +1해준다.

3. 만약 while문을 다 돌았는데도 함수가 리턴되지 않았다면 T_max 시간 내로 모든 소가 공연을 마치지 못한다는 말이다. 그러므로 이때는 T_max+1을 리턴해준다.