[データ構造] 2分ヒープをPythonで実装

Python のリストを使い Max Heap を2分ヒープで実装します。

初期化

配列の最大要素数を定数MAX_NUM_ITEMSとして定めます。

また、self.heap_size という定数を用意して、配列の最後の要素にアクセスする際にインデックスとして使用します。

MAX_NUM_ITEMS = 20

class MaxHeap(object):

    def __init__(self):
        self.heap = [0] * MAX_NUM_ITEMS
        self.heap_size  = 0

値の挿入

配列の一番最後に値を挿入し、 heapify_up を実行します。

heapify_up は、もし親要素の値の方が小さい場合は値を入れ替え、さらに再帰的にheapify_upを行います。

MAX_NUM_ITEMS = 20

class MaxHeap(object):

    def __init__(self):
        self.heap = [0] * MAX_NUM_ITEMS
        self.heap_size  = 0

    # O(log N)
    def insert(self, item):
        if self.heap_size == MAX_NUM_ITEMS:
            raise Exception("ヒープはいっぱい")

        # O(1)
        self.heap[self.heap_size] = item
        self.heap_size += 1

        # O(log N)
        self.heapify_up(self.heap_size-1)

    # O(log N)
    def heapify_up(self, index):
        parent_index = (index-1) // 2
        if index > 0 and self.heap[index] > self.heap[parent_index]:
            self.heap[index], self.heap[parent_index] = self.heap[parent_index], self.heap[index]
            self.heapify_up(parent_index)

最大値

最大値の確認peekは、配列の最初の要素を確認するだけです。

最大値の取り出しextractは、まず配列の最初の要素を最大値とし、その後最初の要素と最後の要素を入れ替え、heap_sizeを減らすことでヒープから最後の要素を外します。

最後にheapify_down を実行します。

heapify_down は自身、左側の子、右側の子、それぞれの要素を比較し、子の方が自身より大きい場合は交換し、再帰的にheapify_downを実行します。

MAX_NUM_ITEMS = 20

class MaxHeap(object):

    def __init__(self):
        self.heap = [0] * MAX_NUM_ITEMS
        self.heap_size  = 0

    # O(log N)
    def insert(self, item):
        if self.heap_size == MAX_NUM_ITEMS:
            raise Exception("ヒープはいっぱい")

        # O(1)
        self.heap[self.heap_size] = item
        self.heap_size += 1

        # O(log N)
        self.heapify_up(self.heap_size-1)

    # O(log N)
    def heapify_up(self, index):
        parent_index = (index-1) // 2
        if index > 0 and self.heap[index] > self.heap[parent_index]:
            self.heap[index], self.heap[parent_index] = self.heap[parent_index], self.heap[index]
            self.heapify_up(parent_index)

    # O(1)
    def peek(self):
        return self.heap[0]

    # O(log N)
    def extract(self):
        max_data = self.heap[0]

        self.heap[0], self.heap[self.heap_size-1] = self.heap[self.heap_size-1], self.heap[0]
        self.heap_size -= 1

        # O(log N)
        self.heapify_down(0)
        return max_data

    # O(log N)
    def heapify_down(self, index):
        index_left_child = 2 * index + 1
        index_right_child = 2 * index + 2

        largest_index = index

        if index_left_child < self.heap_size and self.heap[index_left_child] > self.heap[largest_index]:
            largest_index = index_left_child
        if index_right_child < self.heap_size and self.heap[index_right_child] > self.heap[largest_index]:
            largest_index = index_right_child

        if index != largest_index:
            self.heap[index], self.heap[largest_index] = self.heap[largest_index], self.heap[index]
            self.heapify_down(largest_index)

ヒープソート

最大値の取り出しを要素数だけ行えば、ソートされた配列になります。

MAX_NUM_ITEMS = 8

class MaxHeap(object):

    def __init__(self):
        self.heap = [0] * MAX_NUM_ITEMS
        self.heap_size  = 0

    # O(log N)
    def insert(self, item):
        if self.heap_size == MAX_NUM_ITEMS:
            raise Exception("ヒープはいっぱい")

        # O(1)
        self.heap[self.heap_size] = item
        self.heap_size += 1

        # O(log N)
        self.heapify_up(self.heap_size-1)

    # O(log N)
    def heapify_up(self, index):
        parent_index = (index-1) // 2
        if index > 0 and self.heap[index] > self.heap[parent_index]:
            self.heap[index], self.heap[parent_index] = self.heap[parent_index], self.heap[index]
            self.heapify_up(parent_index)

    # O(1)
    def peek(self):
        return self.heap[0]

    # O(log N)
    def extract(self):
        max_data = self.heap[0]

        self.heap[0], self.heap[self.heap_size-1] = self.heap[self.heap_size-1], self.heap[0]
        self.heap_size -= 1

        # O(log N)
        self.heapify_down(0)
        return max_data

    # O(log N)
    def heapify_down(self, index):
        index_left_child = 2 * index + 1
        index_right_child = 2 * index + 2

        largest_index = index

        if index_left_child < self.heap_size and self.heap[index_left_child] > self.heap[largest_index]:
            largest_index = index_left_child
        if index_right_child < self.heap_size and self.heap[index_right_child] > self.heap[largest_index]:
            largest_index = index_right_child

        if index != largest_index:
            self.heap[index], self.heap[largest_index] = self.heap[largest_index], self.heap[index]
            self.heapify_down(largest_index)

    # O (N log N)
    def heap_sort(self):
        size = self.heap_size
        # O(N) times
        for _ in range(size):
            # O (log N)
            self.extract()
        return self.heap


if __name__ == '__main__':
    maxheap = MaxHeap()
    maxheap.insert(20)
    maxheap.insert(50)
    maxheap.insert(10)
    maxheap.insert(30)
    maxheap.insert(100)
    maxheap.insert(80)
    # 100
    print(maxheap.peek())
    # [100, 50, 80, 20, 30, 10, 0, 0]
    print(maxheap.heap)
    sorted_array = maxheap.heap_sort()
    # [10, 20, 30, 50, 80, 100, 0, 0]
    print(sorted_array)

Python でのヒープ

heapq というモジュールで2分ヒープが提供されており、自分で実装する必要はありません。

heapq — ヒープキューアルゴリズム

pypi で検索すると、フィボナッチヒープも出てきます。