## AVL tree addition and spin Python implementation

Crispy crisp 2020-11-14 16:22:33

``````from collections import deque
from dataStructures.tree.biTree.bst import BiTreeNode, BST
class AVLNode(BiTreeNode):
def __init__(self, data):
BiTreeNode.__init__(self, data)
self.bf = 0 # Survival balance factor
class AVLTree(BST):
def __init__(self, li=None):
BST.__init__(self, li)
def rotate_left(self, p, c):
#c It is the present. node,p yes node.parent
s2 = c.lchild
p.rchild = s2
if s2:
s2.parent = p
c.lchild = p
p.parent = c
p.bf = 0
c.bf = 0
return c
def rotate_right(self, p, c):
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):
g = c.lchild
# First right , It becomes right and right
s3 = g.rchild
c.lchild = s3
if s3:
s3.parent = c
g.rchild = c
c.parent = g
s2 = g.lchild
p.rchild = s2
if s2:
s2.parent = p
g.lchild = p
p.parent = g
# to update bf（ On the right - On the left ）
# （1） Insert s3
if g.bf > 0:
p.bf = -1
c.bf = 0
# （2） Insert s2
elif g.bf < 0:
p.bf = 0
c.bf = 1
else: # What's inserted is g
p.bf = 0
c.bf = 0
# What's inserted is g（ signify s1、s2、s3、s4 It's empty ）
return g
def rotate_left_right(self, p, c):
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
# to update bf
if g.bf < 0: # Insert s2
p.bf = 1
c.bf = 0
elif g.bf > 0: # Insert s3
p.bf = 0
c.bf = -1
else:
p.bf = 0
c.bf = 0
return g
def insert_no_rec(self, val):
# 1. and BST equally , Insert
p = self.root
if not p: # Empty tree
self.root = AVLNode(val)
return
while True:
if val < p.data:
if p.lchild:
p = p.lchild # node Storage is the inserted node
else: # The left child doesn't exist
p.lchild = AVLNode(val)
p.lchild.parent = p
node = p.lchild # node What is stored is the inserted node
break
elif val > p.data:
if p.rchild:
p = p.rchild
else:
p.rchild = AVLNode(val)
p.rchild.parent = p
node = p.rchild
break
else: # val == p.data
return
# 2. to update balance factor
while node.parent: # node.parent Not empty
if node.parent.lchild == node: # The pass is from the left subtree , The left tree is even heavier
# to update node.parent Of bf -= 1
if node.parent.bf < 0: # original node.parent.bf == -1, After updating, it becomes -2
# Make a spin
# see node Which side is heavy
g = node.parent.parent # To connect the subtrees after rotation
x = node.parent # Root of subtree before rotation
# From the current node.bf Decide whether to insert left or right
if node.bf > 0:
n = self.rotate_left_right(node.parent, node) # n Is the head node of the subtree after rotation
else:
n = self.rotate_right(node.parent, node)
# Remember ： hold n and g Connect
elif node.parent.bf > 0: # original node.parent.bf = 1, After updating, it becomes 0
node.parent.bf = 0
break
else: # original node.parent.bf = 0, After updating, it becomes -1
node.parent.bf = -1
node = node.parent
continue
else: # The pass is from the right subtree , The right tree is even heavier
# to update node.parent.bf += 1
if node.parent.bf > 0: # original node.parent.bf == 1, After updating, it becomes 2
# Make a spin
# see node Which side is heavy
g = node.parent.parent # To connect the subtrees after rotation
x = node.parent # Root of subtree before rotation
# From this we can judge whether it is the left subtree or the right subtree of the parent node
if node.bf < 0: # node.bf = -1
# Right left structure
n = self.rotate_right_left(node.parent, node)
else: # node.bf = 1
# Right right structure
n = self.rotate_left(node.parent, node)
# Remember to connect
elif node.parent.bf < 0: # original node.parent.bf = -1, After updating, it becomes 0
node.parent.bf = 0
break
else: # original node.parent.bf = 0, After updating, it becomes 1
node.parent.bf = 1
node = node.parent
continue
# Link the subtree after rotation
n.parent = g
if g: # g Not empty
# Judge whether it was the left child or the right child
if x == g.lchild:
g.lchild = n
else:
g.rchild = n
break
else:
self.root = n
break
def level_order(self, root):
queue = deque()
queue.append(root)
while len(queue) > 0:
node = queue.popleft()
print(node.data, end=',')
if node.lchild:
queue.append(node.lchild)
if node.rchild:
queue.append(node.rchild)
tree = AVLTree([10,5,0])
``````