[백준]11279-최대힙

문제바로가기

제목 그대로 최대힙을 구현하는 문제인데 대부분 라이브러리를 이용하지만 나의 목표는 알고리즘 실력향상이므로 없이 짜보았다.

import sys

heap = []
def heappush(heap, x):
    heap.append(x)
    i = len(heap) - 1
    while i > 0:
        parent = (i - 1)//2
        if heap[parent] < heap[i]:
            heap[i], heap[parent] = heap[parent], heap[i]
            i = parent
        else:
            break
    return heap

def heappop(heap):
    if not heap:
        return None
    
    root = heap[0]
    last = heap.pop()

    if heap:
        heap[0] = last
        i = 0
        while True:
            left = 2 * i + 1
            right = 2 * i + 2
            largest = i

            if left < len(heap) and heap[left] > heap[largest]:
                largest = left
            if right < len(heap) and heap[right] > heap[largest]:
                largest = right

            if largest == i:
                break

            heap[i], heap[largest] = heap[largest], heap[i]
            i = largest

    return root

n = int(sys.stdin.readline())
heap = []
for i in range(n):
    num = int(sys.stdin.readline())
    if num == 0:
        if not heap:
            print(0)
        elif len(heap) > 0:
            print(heappop(heap))
    else:
        heappush(heap, num)

여기서 실수했는데 못찾은 부분이 몇가지 있다. 지피티한테 물어봤다.

1. pop함수에서 return을 heap을 했던 것.
   heap.pop()을 하면 루트값(삭제할 최댓값)을 잃어버리는데, 그걸 리턴하지 않고 리스트를 반환하고 있음.
2. print(heappop(heap)) 대신 print(max(heap))이라고 쓴 것.
   최대힙은 어차피 루트에 최대값이 보장되는데 이를 이용하지 않고 O(n)인 max 연산을 또 했다.

직접 짜보니까 더 도움이 된 듯.