자료구조

[python] minheap 구현하기

유병각 2022. 7. 3. 22:38
class MinHeap:
    def __init__(self):
        self.queue = [None]
    
    def insert(self, num):
        self.queue.append(num)
        lastElementIdx = len(self.queue) - 1
        parentIdx = lastElementIdx // 2
        while (parentIdx >= 1 and self.queue[lastElementIdx] < self.queue[parentIdx]):
            self.queue[parentIdx], self.queue[lastElementIdx] = self.queue[lastElementIdx], self.queue[parentIdx]
            lastElementIdx = parentIdx
            parentIdx = parentIdx // 2
    
    def getLastIdx(self):
        return len(self.queue) - 1
    
    def swap(self, idx1, idx2):
        self.queue[idx1], self.queue[idx2] = self.queue[idx2], self.queue[idx1]
    
    def getLeftChildIdx(self, idx):
        lci = idx * 2
        if (self.getLastIdx() < lci):
            return -1
        else:
            return lci
    def getRightChildIdx(self, idx):
        rci = idx * 2 + 1
        if (self.getLastIdx() < rci):
            return -1
        else:
            return rci
    def getParentIdx(self, idx):
        return idx // 2
    
    def minHeapify(self, idx):
        #반복이 끝나는 조건 1. 현재 노드의 자식이 존재하지 않는다. 2. 현재 노드의 자식들 중 최솟값 보다 현재 노드값이 더 작다
        minIdx = idx  
        lci = self.getLeftChildIdx(idx)
        rci = self.getRightChildIdx(idx)
        
        if (lci != -1 and self.queue[lci] < self.queue[minIdx]):
            minIdx = lci
        if (rci != -1 and self.queue[rci] < self.queue[minIdx]):
            minIdx = rci
        
        if (minIdx != idx):
            self.swap(idx, minIdx)
            self.minHeapify(minIdx)
    
    def pop(self):
        lastIdx = self.getLastIdx()
        if (lastIdx == 0):
            return None
        self.swap(1, lastIdx)
        minV = self.queue.pop()
        self.minHeapify(1)
        return minV