1208. [S/W 문제해결 기본] 1일차 - Flatten

https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AV139KOaABgCFAYh

SW Expert Academy

heapq를 사용하여 max heap, min heap을 구현

각 heapq마다 최대 최소를 구해서 dump 과정 수행

 

import heapq

for t in range(1, 10 + 1) :
    num_dump = int(input())
    boxes = list(map(int, input().split()))

    min_heap = []
    max_heap = []
    for box in boxes :
        heapq.heappush(min_heap, box)
        heapq.heappush(max_heap, -box)

    for n in range(num_dump) :
        minimum = heapq.heappop(min_heap)
        min_max = heapq.nlargest(len(min_heap), min_heap)[1:]
        maximum = -1 * heapq.heappop(max_heap)
        min_max = heapq.nlargest(len(max_heap), min_heap)[1:]

        heapq.heappush(min_heap, minimum + 1)
        heapq.heappush(max_heap, -(maximum - 1))

    ans = -heapq.heappop(max_heap) - heapq.heappop(min_heap)

    print(f'#{t} {ans}')