[백준 11286] 절댓값 힙 with Python

문제 링크

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

11286번: 절댓값 힙

풀이

import sys

class Heap:
    def __init__(self, L=[]):
        self.arr = L

    def compare(self, a, b):
        if abs(self.arr[a]) > abs(self.arr[b]):
            return True
        elif abs(self.arr[a]) < abs(self.arr[b]):
            return False
        else:
            return self.arr[a] > self.arr[b]

    def heapify_down(self, k, n):
        while 2*k+1 < n:
            L, R = 2*k+1, 2*k+2
            if L < n and self.compare(k, L):
                m = L
            else:
                m = k
            if R < n and self.compare(m, R):
                m = R
            if m != k:
                self.arr[k], self.arr[m] = self.arr[m], self.arr[k]
                k = m
            else:
                break

    def heapify_up(self, k):
        while k > 0 and self.compare((k-1)//2, k):
            self.arr[k], self.arr[(k-1)//2] = self.arr[(k-1)//2], self.arr[k]
            k = (k-1)//2

    def insert_key(self, key):
        self.arr.append(key)
        self.heapify_up(len(self.arr)-1)

    def delete_max(self):
        if len(self.arr) == 0: return None
        self.arr[0], self.arr[len(self.arr)-1] = self.arr[len(self.arr)-1], self.arr[0]
        key = self.arr.pop()
        self.heapify_down(0, len(self.arr))
        return key

    def is_empty(self):
        return len(self.arr) == 0

N = int(sys.stdin.readline())
heap = Heap()
output = []
for i in range(N):
    x = int(sys.stdin.readline())
    if x == 0:
        if heap.is_empty():
            output.append(0)
        else:
            output.append(heap.delete_max())
    else:
        heap.insert_key(x)
        
print('\n'.join(map(str, output)))

직접 Heap 클래스를 구현하여 해결하였다.

compare 메서드는 a와 b를 인자로 받아서 arr[a]가 arr[b]보다 절댓값이 큰지,

만약 절댓값이 같다면 arr[a]가 음수인지 여부를 반환한다.

heap에 원소를 push할 때 또는 pop한 뒤에는 arr가 compare 메서드를 기준으로 정렬된다.

 

 

import sys
import heapq

N = int(sys.stdin.readline())
plus_heap = []
minus_heap = []
output = []
for _ in range(N):
    x = int(sys.stdin.readline())
    if x > 0:
        heapq.heappush(plus_heap, x)
    elif x < 0:
        heapq.heappush(minus_heap, -x)
    else:
        if not plus_heap and not minus_heap:
            output.append(0)
        elif not plus_heap:
            output.append(-heapq.heappop(minus_heap))
        elif not minus_heap:
            output.append(heapq.heappop(plus_heap))
        elif plus_heap[0] < minus_heap[0]:
            output.append(heapq.heappop(plus_heap))
        else:
            output.append(-heapq.heappop(minus_heap))

print('\n'.join(map(str, output)))

파이썬의 heapq 모듈을 활용하면 더 쉽게 해결할 수 있다.

heapq는 최소힙

heapq로 문제를 풀기 위해 plus_heap, minus_heap으로 두 개의 최소힙이 필요하다.

minus_heap에 push할 때는 양수로 만들어서 push하고, pop할 때는 pop된 원소를 다시 음수로 만들어줘야 한다.