在Python中實現樹的平衡可以通過實現平衡二叉樹(AVL樹)來實現。AVL樹是一種自平衡二叉搜索樹,即樹的每個節點的左子樹和右子樹的高度差不超過1。
下面是一個簡單的Python實現:
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
使用上述代碼可以創建一個AVL樹,然后通過調用insert
方法向樹中插入新節點,保持樹的平衡。