AlgoDatSoSe24/SoSe24/lec04_trees/b_tree.py

from SoSe24.algodat.foundation import AlgoDatArray, AlgoDatValue, read_int_sequence
from time import perf_counter as pfc


class BTreeNode:
    def __init__(self, m: int):
        self.n = 0
        self.leaf = True
        self.keys = AlgoDatArray(2 * m - 1)
        self.children = [None] * (2 * m)

    def __str__(self):
        return "(" + " ".join([str(self.keys[i]) for i in range(self.n)]) + ")"

class BTree:
    def __init__(self, m: int):
        self.m = m
        self.root = BTreeNode(m)

    def search(self, key: AlgoDatValue, start: BTreeNode = None) -> BTreeNode:
        if not start:
            start = self.root
        i = 0
        while i < start.n and key > start.keys[i]:
            i += 1
        if i < start.n and key == start.keys[i]:
            return start
        if start.leaf:
            return None
        return self.search(key, start.children[i])

    def split_child(self, parent: BTreeNode, i: int):
        child = parent.children[i]
        h = BTreeNode(self.m)
        h.leaf = child.leaf
        h.n = self.m - 1
        for j in range(self.m - 1):
            h.keys[j] = child.keys[j + self.m]
        if not h.leaf:
            for j in range(self.m):
                h.children[j] = child.children[j + self.m]
            for j in range(self.m, child.n + 1):
                child.children[j] = None
        child.n = self.m - 1
        for j in range(parent.n, i, -1):
            parent.children[j + 1] = parent.children[j]
            parent.keys[j] = parent.keys[j - 1]
        parent.children[i + 1] = h
        parent.keys[i] = child.keys[self.m - 1]
        parent.n += 1

    def insert(self, k: AlgoDatValue):
        r = self.root
        if r.n == 2 * self.m - 1:
            h = BTreeNode(self.m)
            self.root = h
            h.leaf = False
            h.n = 0
            h.children[0] = r
            self.split_child(h, 0)
            self.insert_in_node(h, k)
        else:
            self.insert_in_node(r, k)

    def insert_in_node(self, start: BTreeNode, k: AlgoDatValue):
        i = start.n
        if start.leaf:
            while i >= 1 and k < start.keys[i-1]:
                start.keys[i] = start.keys[i-1]
                i -= 1
            start.keys[i] = k
            start.n += 1
        else:
            j = 0
            while j < start.n and k > start.keys[j]:
                j += 1
            if start.children[j].n == 2 * self.m - 1:
                self.split_child(start, j)
                if k > start.keys[j]:
                    j += 1
            self.insert_in_node(start.children[j], k)

    def walk(self, start: BTreeNode = None):
        if not start:
            start = self.root
        i = 0
        while i < start.n:
            if not start.leaf:
                self.walk(start.children[i])
            print(start.keys[i], end=" ")
            i += 1
        if not start.leaf:
            self.walk(start.children[i])

    def height(self, start: BTreeNode = None):
        if not start:
            start = self.root
        if start.leaf:
            return 0
        return 1 + self.height(start.children[0])

    def graph_walk(self):
        queue = [ self.root ]
        with open("../../btree.gv", "w") as file:
            file.write("digraph BTree {\n")
            file.write("    node [fontname=\"Arial\"];\n")
            while queue:
                current = queue.pop(0)
                p = str(current)
                file.write(f'"{p}"; \n')
                i = 0
                while i <= current.n:
                    if not current.leaf:
                        queue.append(current.children[i])
                        c = str(current.children[i])
                        file.write(f'"{p}" -> "{c}";\n')
                    i += 1
            file.write("}")


if __name__ == "__main__":
    z = read_int_sequence("../../seq2.txt")
    start = pfc()
    tree = BTree(3)
    for i in z:
        tree.insert(i)
    print(f"Height: {tree.height()}")
    tree.walk()
    tree.graph_walk()
    s = tree.search(AlgoDatValue(0))
    print(f"\nKnoten mit 0: {str(s)}")
    print(f"Dauer: {pfc() - start:.4f}s")
    AlgoDatValue.summary()