python平衡树

参考：www.bilibili.com/video/BV1E4…

python实现

class Node:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

    def __str__(self):
        return f"node [value={self.value}]"

    def preOrder(self, node):
        if node is None:
            return
        else:
            print(node)
            if node.left:
                node.left.preOrder(node.left)
            if node.right:
                node.right.preOrder(node.right)

    def addNode(self, node):
        if node is None or not isinstance(node, Node):
            return
        if self.value > node.value:
            if self.left is None:
                self.left = node
            else:
                self.left.addNode(node)
        else:
            if self.right is None:
                self.right = node
            else:
                self.right.addNode(node)

    def searchNode(self, value):
        if self.value == value:
            return self
        else:
            if self.left:
                node = self.left.searchNode(value)
                if node:
                    return node
            if self.right:
                node = self.right.searchNode(value)
                if node:
                    return node

    def getParentNode(self, value):
        if (self.left and self.left.value == value) or (self.right and self.right.value == value):
            return self
        # 向左递归
        if self.value > value:
            node = self.left.getParentNode(value)
            if node:
                return node
        # 向右递归
        elif self.value < value:
            node = self.right.getParentNode(value)
            if node:
                return node
        else:
            return None

    # 左树高度
    def getLeftHeight(self):
        if not self.left:
            return 0
        return self.left.treeHeight()

    # 右树高度
    def getRightHeight(self):
        if not self.right:
            return 0
        return self.right.treeHeight()

    # 树的高度
    def treeHeight(self):
        return max(0 if not self.left else self.left.treeHeight(), 0 if not self.right else self.right.treeHeight()) + 1


class AVLTree:
    def __init__(self, root=None):
        self.root = root

    # 左旋
    def _leftRotate(self):
        if self.hasRootNode():
            print("空树")
            return
        # 1.创建新的节点
        newNode = Node(self.root.value)
        # 2.newNode的左节点指向root节点的左节点
        newNode.left = self.root.left
        # 3.newNode的右节点指向root节点的右子树的左节点
        newNode.right = self.root.right.left
        # 4.root节点的值换成右子树的值
        self.root.value = self.root.right.value
        # 5.root节点的左节点指向newNode
        self.root.left = newNode
        # 6.root节点的右节点指向原来root节点右子树的右子树
        self.root.right = self.root.right.right

    # 右旋
    def _rightRotate(self):
        if self.hasRootNode():
            print("空树")
            return
        # 1.创建新的节点
        newNode = Node(self.root.value)
        # 2.newNode的右节点指向root的右子树
        newNode.right = self.root.right
        # 3.newNode的左节点指向root左子树的右子树
        newNode.left = self.root.left.right
        # 4.root的值替换成左子树的值
        self.root.value = self.root.left.value
        # 5.root节点的左边指向左子树的左子树
        self.root.left = self.root.left.left
        # 6.root节点右子树指向新节点
        self.root.right = newNode

    # 添加节点
    def addNode(self, node):
        if node is None or not isinstance(node, Node):
            return
        # 没有根节点的情况
        if self.hasRootNode():
            self.root = node
        # 有根节点的情况
        else:
            self.root.addNode(node)

        # 判断是否左旋
        if self.root.getRightHeight() - self.root.getLeftHeight() > 1:
            if self.root.right and self.root.right.getLeftHeight() - self.root.right.getRightHeight() > 0:
                tempNode = self.root
                self.root = self.root.right
                self._rightRotate()
                self.root = tempNode
            self._leftRotate()
            return
        # 判断是否右旋
        if self.root.getLeftHeight() - self.root.getRightHeight() > 1:
            # root节点的左子树的右子树高度大于左子树
            if self.root.left and self.root.left.getRightHeight() - self.root.left.getLeftHeight() > 0:
                tempNode = self.root
                self.root = self.root.left
                self._leftRotate()
                self.root = tempNode
            self._rightRotate()

    # 查找
    def searchNode(self, value):
        if not self.hasRootNode():
            return self.root.searchNode(value)
        else:
            print("当前树为空")

    # 前序遍历
    def preOrder(self):
        if not self.hasRootNode():
            self.root.preOrder(self.root)

    # 判断root是否为空
    def hasRootNode(self):
        return self.root is None

    def getMinValue(self, node):
        if node:
            while node.left:
                node = node.left
            self.delNode(node.value)
            return node.value
        else:
            print("node为空节点")

    # 3种情况
    # 1.删除的是叶子节点，也就是左右节点都为空
    # 2.删除的是有一个叶子节点的子节点
    # 3.删除的是有2个叶子节点的子节点
    def delNode(self, value):
        if self.hasRootNode():
            print("树为空")
        else:
            targetNode = self.searchNode(value)
            # 没找到目标节点
            if not targetNode:
                print("没有找到该节点")
                return
            # 如果没有找到早就return了，不会进去下面的步骤
            if not self.root.left and not self.root.right:
                self.root = None
                return
            parentNode = self.getParent(value)
            # 叶子节点
            if not targetNode.left and not targetNode.right:
                if parentNode.left and parentNode.left.value == targetNode.value:
                    parentNode.left = None
                elif parentNode.right and parentNode.right.value == targetNode.value:
                    parentNode.right = None
            # 有左右子节点
            elif targetNode.left and targetNode.right:
                targetNode.value = self.getMinValue(targetNode.right)
            # 有左或右一个子节点
            else:
                # 目标节点的左子节点不为空
                if targetNode.left:
                    # 父节点是否为空
                    if parentNode:
                        # 目标节点是父节点的左节点
                        if parentNode.left and parentNode.left.value == value:
                            parentNode.left = targetNode.left
                        else:
                            parentNode.right = targetNode.left
                    else:
                        self.root = targetNode.left
                # 目标节点的右子节点不为空
                else:
                    # 父节点是否为空
                    if parentNode:
                        if parentNode.left and parentNode.left.value == value:
                            parentNode.left = targetNode.right
                        else:
                            parentNode.right = targetNode.right
                    else:
                        self.root = targetNode.right

    def getParent(self, value):
        if self.hasRootNode():
            print("树为空")
        else:
            return self.root.getParentNode(value)


avlTree = AVLTree(Node(10))
avlTree.addNode(Node(5))
avlTree.addNode(Node(20))
avlTree.addNode(Node(11))
avlTree.addNode(Node(25))
avlTree.addNode(Node(15))
print(avlTree.root.treeHeight())
print(avlTree.root.getLeftHeight())
print(avlTree.root.getRightHeight())

avlTree.preOrder()