FibHeap.py

import heapq
import random


class FibNode:
    def __init__(self, key):
        self.degree = 0
        self.parent = None
        self.children = []
        self.key = key
        self.right = self
        self.left = self


class FibHeap:
    def __init__(self):
        self.min = None
        self.size = 0

    def insert(self, key):
        x = FibNode(key)
        if not self.min:
            self.min = x
        else:
            self.add_root(x)
            if x.key < self.min.key:
                self.min = x
        self.size += 1

    def add_root(self, x: FibNode):
        """
        根の連結リストにノードを追加する
        """
        if not self.min:
            self.min = x
            return

        left = self.min.left
        left.right = x
        self.min.left = x
        x.right = self.min
        x.left = left

    def extract_min(self):
        """
        最小値を取り出す
        """
        z = self.min
        if not z:
            return None
        for child in z.children:
            self.add_root(child)
            child.parent = None

        self.remove_root(z)
        if z == z.right:
            # これが最後のノードならヒープを空にする
            self.min = None
        else:
            # とりあえず適当に先頭を決める
            self.min = z.right
            self.consolidate()
        self.size -= 1
        return z.key

    def consolidate(self):
        # 根の連結リストを作る
        root_list = [self.min]
        while root_list[len(root_list) - 1].right != root_list[0]:
            last_right = root_list[len(root_list) - 1].right
            root_list.append(last_right)

        root_of_degree = [None] * (self.size + 1)
        for parent in root_list:
            degree = parent.degree
            while root_of_degree[degree]:
                child = root_of_degree[degree]
                if parent.key > child.key:
                    t = parent
                    parent = child
                    child = t
                self.link(child, parent)
                root_of_degree[degree] = None
                degree += 1
            root_of_degree[degree] = parent

        self.min = None
        for node in root_of_degree:
            if not node:
                continue

            if not self.min:
                node.right = node
                node.left = node
                self.min = node
            else:
                self.add_root(node)
                if node.key < self.min.key:
                    self.min = node

    def link(self, child: FibNode, parent: FibNode):
        """
        子ノードを親ノードの下に接続する
        """
        self.remove_root(child)
        parent.degree += 1
        parent.children.append(child)
        child.parent = parent

    def remove_root(self, x: FibNode):
        """
        根の連結リストからノードを削除する
        """
        r = x.right
        l = x.left
        r.left = l
        l.right = r


def main():
    fib_heap = FibHeap()
    heap = []
    for _ in range(100000):
        if len(heap) > 0 and random.randint(1, 2) == 1:
            if heapq.heappop(heap) != fib_heap.extract_min():
                raise ValueError()
        else:
            x = random.randint(0, 100000)
            heapq.heappush(heap, x)
            fib_heap.insert(x)


if __name__ == '__main__':
    main()