用Python实现树形结构的AVL-RT树

class Node:
    def __init__(self, key=None, data=None, left=None, right=None, height=0, bbox=None):
        self.key = key
        self.data = data
        self.left = left
        self.right = right
        self.height = height
        self.bbox = bbox

    def __repr__(self):
        return f"<Node key={self.key}, data={self.data}, bbox={self.bbox}>"

    def __str__(self):
        return f"<Node key={self.key}, data={self.data}, bbox={self.bbox}>"

class AVL_RT:
    def __init__(self, bbox=None):
        self.root = None
        self.dim = len(bbox) if bbox is not None else 2

    def height(self, node):
        if node is None:
            return -1
        else:
            return node.height

    def balance_factor(self, node):
        return self.height(node.left) - self.height(node.right)

    def insert(self, key, data=None, bbox=None):
        new_node = Node(key=key, data=data, bbox=bbox)
        if self.root is None:
            self.root = new_node
        else:
            self.root = self._insert(self.root, new_node)

    def _insert(self, node, new_node):
        if node is None:
            return new_node
        elif new_node.key < node.key:
            node.left = self._insert(node.left, new_node)
        else:
            node.right = self._insert(node.right, new_node)
        node.height = max(self.height(node.left), self.height(node.right)) + 1
        bf = self.balance_factor(node)
        if bf > 1 and new_node.key < node.left.key:
            return self._rotate_right(node)
        if bf < -1 and new_node.key > node.right.key:
            return self._rotate_left(node)
        if bf > 1 and new_node.key > node.left.key:
            node.left = self._rotate_left(node.left)
            return self._rotate_right(node)
        if bf < -1 and new_node.key < node.right.key:
            node.right = self._rotate_right(node.right)
            return self._rotate_left(node)
        return node

    def _rotate_right(self, node):
        left = node.left
        left_right = left.right
        left.right = node
        node.left = left_right
        node.height = max(self.height(node.left), self.height(node.right)) + 1
        left.height = max(self.height(left.left), self.height(left.right)) + 1
        return left

    def _rotate_left(self, node):
        right = node.right
        right_left = right.left
        right.left = node
        node.right = right_left
        node.height = max(self.height(node.left), self.height(node.right)) + 1
        right.height = max(self.height(right.left), self.height(right.right)) + 1
        return right

    def search(self, bbox):
        res = []
        self._search(self.root, bbox, res)
        return res

    def _search(self, node, bbox, res):
        if node is None:
            return
        if node.bbox.intersects(bbox):
            res.append(node.data)
        if node.left and node.left.bbox.intersects(bbox):
            self._search(node.left, bbox, res)
        if node.right and node.right.bbox.intersects(bbox):
            self._search(node.right, bbox, res)

if __name__ == '__main__':
    tree = AVL_RT([(-10, 10), (-10, 10)])
    tree.insert(10, "pidancode.com", [(-1, 1), (-1, 1)])
    tree.insert(5, "皮蛋编程", [(0, 2), (-3, -1)])
    tree.insert(-5, "Python", [(-5, -3), (4, 6)])
    res = tree.search([(-2, 2), (-2, 2)])
    print(res)