백준 1927번 최소힙 파이썬

문제

널리 잘 알려진 자료구조 중 최소 힙이라는 것이 있다. 최소 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성하시오.

1. 배열에 자연수 x를 넣는다.
2. 배열에서 가장 작은 값을 출력하고, 그 값을 배열에서 제거한다.

프로그램은 처음에 비어있는 배열에서 시작하게 된다.

입력

첫째 줄에 연산의 개수 N(1≤N≤100,000)이 주어진다. 다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다. 만약 x가 자연수라면 배열에 x라는 값을 넣는(추가하는) 연산이고, x가 0이라면 배열에서 가장 작은 값을 출력하고 그 값을 배열에서 제거하는 경우이다. 입력되는 자연수는 2^31보다 작다.

출력

입력에서 0이 주어진 회수만큼 답을 출력한다. 만약 배열이 비어 있는 경우인데 가장 작은 값을 출력하라고 한 경우에는 0을 출력하면 된다.

 

---

문제 풀이 방법

• 최소 힙의 기본적인 기능을 구현한다.
• 파이썬의 heapq 라이브러리를 이용하면 간단히 힙을 구현할 수 있다.

 

---

import heapq
n = int(input())
heap=[]
result=[]

for _ in range(n):
	data=int(input())
    if data ==0:
    	if heap:
        	result.append(heapq.heappop(heap))
        else:
        	result.append(0)
    else:
    	heapq.heappush(heap, data)
 for data in result:
 	print(data)

class Heap:
    def __init__(self, data):
        self.heap_array = list()
        self.heap_array.append(None)
        self.heap_array.append(data)


    def move_up(self, inserted_idx):
        if inserted_idx<=1:
            return False
        parent_idx = inserted_idx //2
        if self.heap_array[inserted_idx]< self.heap_array[parent_idx]:
            return True
        else:
            return False

    def insert(self, data):
        if len(self.heap_array) ==0:
            self.heap_array.append(None)
            self.heap_array.append(data)
            return True
        self.heap_array.append(data)

        inserted_idx = len(self.heap_array)-1
        while  self.move_up(inserted_idx): # 부모노드와의 비교
            parent_idx = inserted_idx//2
            self.heap_array[inserted_idx] , self.heap_array[parent_idx]=self.heap_array[parent_idx],self.heap_array[inserted_idx]
            inserted_idx=parent_idx

        return True

    def move_down(self, popped_idx):
        left_child_popped_idx = popped_idx*2
        right_child_popped_idx = popped_idx*2+1
        # case 1: 왼쪽 자식 노드도 없을때
        if left_child_popped_idx >=len(self.heap_array):
            return False
        # case 2: 오른쪽 자식 노드만 없을때
        elif right_child_popped_idx >=len(self.heap_array):
            if self.heap_array[popped_idx] > self.heap_array[left_child_popped_idx]:
                return True
            else:
                return False
        # case 3: 왼쪽, 오르쪽 자식 모두 있을때
        else:
            # 왼쪽이 더 큰경우
            if  self.heap_array[left_child_popped_idx]< self.heap_array[right_child_popped_idx]:
                if self.heap_array[popped_idx] > self.heap_array[left_child_popped_idx]:
                    return True
                else:
                    return False
            else:
            # 오른쪽이 더 큰 경우  
                if self.heap_array[popped_idx] > self.heap_array[right_child_popped_idx]:
                    return True
                else:
                    return False
    def pop(self):
        # 원리: 맥스힙의  pop 은 root 를 삭제하는 것이다.
        # 구현 원리는 맨 최상단 root 를 tail 위치의 노드를 복사해 넣고, 마지막 노드를 삭제한 후, 정렬을 해주면 된다.
        if len(self.heap_array)<=1:  # pop 할것이 없는 루트 노드
            print(0)
            return None

        returned_data = self.heap_array[1] 
        self.heap_array[1] =self.heap_array[-1]
        del self.heap_array[-1]
        popped_idx =1
        while  self.move_down( popped_idx):
            left_child_popped_idx = popped_idx*2
            right_child_popped_idx = popped_idx*2+1
            # case 1: 왼쪽 자식 노드도 없을때
            if left_child_popped_idx >=len(self.heap_array):
                pass # 신경 안쓰면 됨
            # case 2: 오른쪽 자식 노드만 없을때
            elif right_child_popped_idx >=len(self.heap_array):
                if self.heap_array[popped_idx] > self.heap_array[left_child_popped_idx]:
                    self.heap_array[popped_idx],self.heap_array[left_child_popped_idx]=self.heap_array[left_child_popped_idx],self.heap_array[popped_idx]
                    popped_idx=left_child_popped_idx
            # case 3: 왼쪽, 오르쪽 자식 모두 있을때
            else:
                # 왼쪽이 더 큰경우
                if  self.heap_array[left_child_popped_idx]<  self.heap_array[right_child_popped_idx]:
                    if self.heap_array[popped_idx] > self.heap_array[left_child_popped_idx]:
                        self.heap_array[popped_idx],self.heap_array[left_child_popped_idx]=self.heap_array[left_child_popped_idx],self.heap_array[popped_idx]
                        popped_idx=left_child_popped_idx
                else:
                # 오른쪽이 더 큰 경우  
                    if self.heap_array[popped_idx] > self.heap_array[right_child_popped_idx]:
                        self.heap_array[popped_idx],self.heap_array[right_child_popped_idx]=self.heap_array[right_child_popped_idx],self.heap_array[popped_idx]
                        popped_idx=right_child_popped_idx
        print(returned_data)
        return returned_data
    
    
N = int(input())
heap =-1
for i in range(N):
    a=int(input())
    
    if heap==-1:
       # print('[]')
        if a==0:
            print(0)
        else:
            heap = Heap(a)
    else:
        
        if a==0:
            heap.pop()
        else:
            heap.insert(a)
       # print(heap.heap_array)