from 2.13_BST import BiTreeNode, BST
AVL树节点类 (继承二叉树节点类)
class AVLNode(BiTreeNode):
def __init__(self, data):
BiTreeNode.__init__(self, data)
# 需要对每个节点跟踪"平衡因子 (balance factor)"参数
self.bf = 0
AVL树类 (继承BST类)
class AVLTree(BST):
def __init__(self, li=None):
BST.__init__(self, li)
def rotate_left(self, p, c):
"""插入p右孩子的右子树 - 左旋"""
# 把 s2 从 c 的左孩子,变为 p 的右孩子
s2 = c.lchild
p.rchild = s2
if s2:
s2.parent = p
# 把 c 是 p 的右子树,变为 p 是 c 的左子树
c.lchild = p
p.parent = c
# 更新 balance factor
p.bf = 0
c.bf = 0
return c
def rotate_right(self, p, c):
"""插入p左孩子的左子树 - 右旋"""
s2 = c.rchild
p.lchild = s2
if s2:
s2.parent = p
c.rchild = p
p.parent = c
p.bf = 0
c.bf = 0
return c
def rotate_right_left(self, p, c):
"""插入p右孩子的左子树 - 右旋+左旋"""
g = c.lchild
# 把 s3 从 g 的右孩子,变为 c 的左孩子
s3 = g.rchild
c.lchild = s3
if s3:
s3.parent = c
# 把 g 是 c 的左子树,变为 c 是 g 的右子树
g.rchild = c
c.parent = g
# 把 s2 从 g 的左孩子,变为 p 的右孩子
s2 = g.lchild
p.rchild = s2
if s2:
s2.parent = p
# 把 g 是 p 的右子树,变为 p 是 g 的左子树
g.lchild = p
p.parent = g
# 更新 balance factor
# 情况1:g 左h-1 右h
if g.bf > 0:
p.bf = -1
c.bf = 0
# 情况2:g 左h 右h-1
elif g.bf < 0:
p.bf = 0
c.bf = 1
# 情况3:插入的就是 g (s1/s2/s3/s4 都是 None)
else:
p.bf = 0
c.bf = 0
g.bf = 0
return g
def rotate_left_right(self, p, c):
"""插入p左孩子的右子树 - 左旋+右旋"""
g = c.rchild
s2 = g.lchild
c.rchild = s2
if s2:
s2.parent = c
g.lchild = c
c.parent = g
s3 = g.rchild
p.lchild = s3
if s3:
s3.parent = p
g.rchild = p
p.parent = g
if g.bf < 0:
p.bf = 1
c.bf = 0
elif g.bf > 0:
p.bf = 0
c.bf = -1
else:
p.bf = 0
c.bf = 0
g.bf = 0
return g
def insert_no_rec(self, val):
# 1. 和 BST 一样插入
p = self.root
if not p:
self.root = AVLNode(val)
return
while True:
if val < p.data:
if p.lchild:
p = p.lchild
else:
p.lchild = AVLNode(val)
p.lchild.parent = p
# node 存储插入的节点
node = p.lchild
break
elif val > p.data:
if p.rchild:
p = p.rchild
else:
p.rchild = AVLNode(val)
p.rchild.parent = p
node = p.rchild
break
# val == p.data
else:
return
# 2. 更新 balance factor - 一直到 root (node.parent not None)
while node.parent:
# 传递是从左子树来的,左子树更沉了
if node == node.parent.lchild:
# 更新 node.parent 的 bf -= 1
# 情况1:node.parent == -1,更新后变成 -2
if node.parent.bf < 0:
# g: 连接旋转之后的子树; x: 旋转前的子树的根
g = node.parent.parent
x = node.parent
# 看 node 哪边沉,确定旋转方式
if node.bf > 0:
n = self.rotate_left_right(node.parent, node)
else:
n = self.rotate_right(node.parent, node)
# 情况2:node.parent == 1,更新后变成 0
elif node.parent.bf > 0:
node.parent.bf = 0
break
# 情况3:node.parent == 0,更新后变成 -1
else:
node.parent.bf = -1
node = node.parent
continue
# 传递是从右子树来的,右子树更沉了
else:
# 更新 node.parent 的 bf += 1
# 情况1:node.parent == 1,更新后变成 2
if node.parent.bf > 0:
g = node.parent.parent
x = node.parent
if node.bf < 0:
n = self.rotate_right_left(node.parent, node)
else:
n = self.rotate_left(node.parent, node)
# 情况2:node.parent == -1,更新后变成 0
elif node.parent.bf < 0:
node.parent.bf = 0
break
# 情况3:node.parent == 0,更新后变成 1
else:
node.parent.bf = 1
node = node.parent
continue
# 连接旋转后的子树
n.parent = g
if g:
if x == g.lchild:
g.lchild = n
else:
g.rchild = n
# 是否继续取决于调整完的 root 的 bf (都是 0,直接 break)
break
else:
self.root = n
break