溫馨提示×

Python TreeNode如何實現(xiàn)樹的平衡

小樊
83
2024-08-17 16:24:40
欄目: 編程語言

在Python中實現(xiàn)樹的平衡可以通過實現(xiàn)平衡二叉樹（AVL樹）來實現(xiàn)。AVL樹是一種自平衡二叉搜索樹，即樹的每個節(jié)點的左子樹和右子樹的高度差不超過1。

下面是一個簡單的Python實現(xiàn)：

class TreeNode:
    def __init__(self, key):
        self.key = key
        self.left = None
        self.right = None
        self.height = 1

class AVLTree:
    def getHeight(self, node):
        if not node:
            return 0
        return node.height

    def getBalance(self, node):
        if not node:
            return 0
        return self.getHeight(node.left) - self.getHeight(node.right)

    def rotateRight(self, y):
        x = y.left
        T = x.right

        x.right = y
        y.left = T

        y.height = 1 + max(self.getHeight(y.left), self.getHeight(y.right))
        x.height = 1 + max(self.getHeight(x.left), self.getHeight(x.right))

        return x

    def rotateLeft(self, x):
        y = x.right
        T = y.left

        y.left = x
        x.right = T

        x.height = 1 + max(self.getHeight(x.left), self.getHeight(x.right))
        y.height = 1 + max(self.getHeight(y.left), self.getHeight(y.right))

        return y

    def insert(self, root, key):
        if not root:
            return TreeNode(key)
        elif key < root.key:
            root.left = self.insert(root.left, key)
        else:
            root.right = self.insert(root.right, key)

        root.height = 1 + max(self.getHeight(root.left), self.getHeight(root.right))

        balance = self.getBalance(root)

        if balance > 1 and key < root.left.key:
            return self.rotateRight(root)
        
        if balance < -1 and key > root.right.key:
            return self.rotateLeft(root)
        
        if balance > 1 and key > root.left.key:
            root.left = self.rotateLeft(root.left)
            return self.rotateRight(root)
        
        if balance < -1 and key < root.right.key:
            root.right = self.rotateRight(root.right)
            return self.rotateLeft(root)

        return root

使用上述代碼可以創(chuàng)建一個AVL樹，然后通過調(diào)用insert方法向樹中插入新節(jié)點，保持樹的平衡。

0 踩