木構造とヒープ

trees_and_heaps.py

# Binary search tree (insert/search) and a binary min-heap
class Node:
    def __init__(self, key):
        self.key = key
        self.left = None
        self.right = None


def bst_insert(root, key):
    if root is None:
        return Node(key)
    if key < root.key:
        root.left = bst_insert(root.left, key)
    elif key > root.key:
        root.right = bst_insert(root.right, key)
    return root  # duplicates are ignored


def bst_search(root, key):
    while root is not None and root.key != key:
        root = root.left if key < root.key else root.right
    return root is not None


def inorder(root, out):
    if root is not None:
        inorder(root.left, out)
        out.append(root.key)
        inorder(root.right, out)
    return out


class MinHeap:
    def __init__(self):
        self.a = []

    def push(self, x):
        a = self.a
        a.append(x)
        i = len(a) - 1
        while i > 0 and a[(i - 1) // 2] > a[i]:  # sift up
            p = (i - 1) // 2
            a[i], a[p] = a[p], a[i]
            i = p

    def pop(self):
        a = self.a
        top, last = a[0], a.pop()
        if a:
            a[0] = last
            i, n = 0, len(a)
            while True:  # sift down
                small = i
                for c in (2 * i + 1, 2 * i + 2):
                    if c < n and a[c] < a[small]:
                        small = c
                if small == i:
                    break
                a[i], a[small] = a[small], a[i]
                i = small
        return top

    def __len__(self):
        return len(self.a)


root = None
for k in [50, 30, 70, 20, 40, 60, 80]:
    root = bst_insert(root, k)
print(inorder(root, []))                     # [20, 30, 40, 50, 60, 70, 80]
print(bst_search(root, 60), bst_search(root, 65))  # True False

heap = MinHeap()
for x in [5, 3, 8, 1, 9, 2]:
    heap.push(x)
print([heap.pop() for _ in range(len(heap))])  # [1, 2, 3, 5, 8, 9]