# coding=utf-8
# AVL树的Python实现(树的节点中包含了指向父节点的指针)
def get_height(node):
return node.height if node else -1
def get_maximum(node):
temp_node = node
while temp_node.right:
temp_node = temp_node.right
return temp_node
def get_minimum(node):
temp_node = node
while temp_node.left:
temp_node = temp_node.left
return temp_node
def preorder_tree_walk(node):
if node:
print node.key, node.height
preorder_tree_walk(node.left)
preorder_tree_walk(node.right)
def left_left_rotate(tree, node):
# 先将 node 和 node_left 之间及其左右节点赋值 (node_left.left node.right 保持不变)
node_left = node.left
node.left = node_left.right
node_left.right = node
if not node.p:
tree.root = node_left
node_left.p = None
elif node == node.p.left:
node.p.left = node_left
node_left.p = node.p
elif node == node.p.right:
node.p.right = node_left
node_left.p = node.p
node.p = node_left
while node:
node.height = max(get_height(node.left), get_height(node.right)) + 1
node = node.p
def right_right_rotate(tree, node):
node_right = node.right
node.right = node_right.left
node_right.left = node
if not node.p:
tree.root = node_right
node_right.p = None
elif node == node.p.left:
node.p.left = node_right
node_right.p = node.p
elif node == node.p.right:
node.p.right = node_right
node_right.p = node.p
node.p = node_right
while node:
node.height = max(get_height(node.left), get_height(node.right)) + 1
node = node.p
def left_right_rotate(tree, node):
right_right_rotate(tree, node.left)
left_left_rotate(tree, node)
def right_left_rotate(tree, node):
left_left_rotate(tree, node.right)
right_right_rotate(tree, node)
class AVLTreeNode(object):
def __init__(self, key):
self.key = key
self.p = None
self.left = None
self.right = None
self.height = 0
class AVLTree(object):
def __init__(self):
self.root = None
def search(self, key):
if not self.root:
return None
else:
return self._search(key)
def _search(self, key):
start = self.root
while start:
if key < start.key:
start = start.left
elif key > start.key:
start = start.right
else:
return start
return None
def insert(self, node):
temp_root = self.root
temp_node = None
# 找到要插入的父节点(temp_node)
while temp_root:
temp_node = temp_root
if node.key < temp_node.key:
temp_root = temp_root.left
elif node.key > temp_node.key:
temp_root = temp_root.right
else:
raise KeyError, "Error!"
# 如果父节点为空 则说明这是一个空树 把 root 赋值即可
if not temp_node:
self.root = node
elif node.key < temp_node.key:
temp_node.left = node
node.p = temp_node
temp_node.height = max(get_height(temp_node.left), get_height(temp_node.right)) + 1
temp_p = temp_node.p
while temp_p:
temp_p.height = max(get_height(temp_p.left), get_height(temp_p.right)) + 1
temp_p = temp_p.p
elif node.key > temp_node.key:
temp_node.right = node
node.p = temp_node
temp_node.height = max(get_height(temp_node.left), get_height(temp_node.right)) + 1
temp_p = temp_node.p
while temp_p:
temp_p.height = max(get_height(temp_p.left), get_height(temp_p.right)) + 1
temp_p = temp_p.p
self.fixup(node)
def fixup(self, node):
if node == self.root:
return
while node:
if get_height(node.left) - get_height(node.right) == 2:
if node.left.left:
left_left_rotate(self, node)
else:
left_right_rotate(self, node)
break
elif get_height(node.right) - get_height(node.left) == 2:
if node.right.right:
right_right_rotate(self, node)
else:
right_left_rotate(self, node)
break
node = node.p
def delete(self, key):
temp_node = self.root
while temp_node:
if key > temp_node.key:
temp_node = temp_node.right
elif key < temp_node.key:
temp_node = temp_node.left
else:
break
if not temp_node:
return False
elif temp_node.left and temp_node.right:
if get_height(temp_node.left) > get_height(temp_node.right):
# 注意删除的时候不是直接把左右子树往上提 而是分别找到左右子树中的最大值和最小值往上提
# 由于是最大子节点 故一定没有右子
node_max = get_maximum(temp_node.left)
if node_max.left:
node_max_p = node_max.p
node_max_p.right = node_max.left
node_max.left.p = node_max_p
node_max.right = temp_node.right
temp_node.right.p = node_max
node_max.left = temp_node.left
temp_node.left.p = node_max
if temp_node.p:
if temp_node == temp_node.p.left:
temp_node.p.left = node_max
node_max.p = temp_node.p
else:
temp_node.p.right = node_max
node_max.p = temp_node.p
else:
self.root = node_max
node_max.p = None
temp_node = node_max
else:
node_min = get_minimum(temp_node.right)
if node_min.right:
node_min_p = node_min.p
node_min_p.left = node_min.right
node_min.right.p = node_min_p
node_min.left = temp_node.left
temp_node.left.p = node_min
node_min.right = temp_node.right
temp_node.right.p = node_min
if temp_node.p:
if temp_node == temp_node.p.left:
temp_node.p.left = node_min
node_min.p = temp_node.p
else:
temp_node.p.right = node_min
node_min.p = temp_node.p
else:
self.root = node_min
node_min.p = None
temp_node = node_min
temp_node.height = max(get_height(temp_node.left), get_height(temp_node.right)) + 1
self.fixup(temp_node)
def main():
number_list = (7, 4, 1, 8, 5, 2, 9, 6, 3)
tree = AVLTree()
for number in number_list:
node = AVLTreeNode(number)
tree.insert(node)
preorder_tree_walk(tree.root)
tree.delete(4)
print '=========='
preorder_tree_walk(tree.root)
if __name__ == '__main__':
main()
可以比较一下上一篇中实现方法的不同
End
