더북(TheBook)

arr 값이 달라질 때마다 pos도 업데이트하면 됩니다. 예를 들어 self.arr[cur_idx]=self.arr[self.parent(cur_idx)]를 보면 cur_idx 위치에 부모 노드를 저장하고 있습니다. arr 내 배치가 바뀐 것이겠죠. 그렇다면 다음에 반드시 self.pos[self.arr[cur_idx].v]=cur_idx 코드로 pos에 수정된 정점의 위치를 업데이트해 주어야 합니다.

코드를 볼까요?

코드 13-10

    def is_empty(self):
        if self.heapsize == 0:
            return True
        return False

    def is_full(self):
        if self.heapsize >= self.MAX_ELEMENTS:
            return True
        return False

    def parent(self, idx):
        return idx >> 1

    def left(self, idx):
        return idx << 1

    def right(self, idx):
        return (idx << 1) + 1

    def push(self, item):
        if self.is_full():
            raise IndexError("the heap is full!!")

        self.heapsize += 1
        cur_idx = self.heapsize

        while cur_idx != 1 and item.w < self.arr[self.parent(cur_idx)].w:
            self.arr[cur_idx] = self.arr[self.parent(cur_idx)]
            # pos 인덱스는 정점, arr은 weight를 키로 만든 최소 힙
            self.pos[self.arr[cur_idx].v] = cur_idx

            cur_idx = self.parent(cur_idx)

        self.arr[cur_idx] = item
        self.pos[item.v] = cur_idx

    def pop(self):
        if self.is_empty():
            return None
        rem_elem = self.arr[1]

        temp = self.arr[self.heapsize]
        self.heapsize -= 1

        cur_idx = 1
        child = self.left(cur_idx)

        while child <= self.heapsize:
            if child < self.heapsize and \
                self.arr[self.left(cur_idx)].w > self.arr[self.right(cur_idx)].w:
                child = self.right(cur_idx)

            if temp.w <= self.arr[child].w:
                break

            self.arr[cur_idx] = self.arr[child]
            self.pos[self.arr[cur_idx].v] = cur_idx

            cur_idx = child
            child = self.left(cur_idx)

        self.arr[cur_idx] = temp
        self.pos[temp.v] = cur_idx

        return rem_elem

코드는 길지만 최소 힙을 충분히 알고 있다면 어렵지 않습니다. arr이 변경될 때 pos도 바꾸어 주어야 한다는 점만 알면 충분합니다.

신간 소식 구독하기
뉴스레터에 가입하시고 이메일로 신간 소식을 받아 보세요.