class Node(OrderedMap.Node):
def __init__(self, element, parent=None, left=None, right=None):
super().__init__(element,parent,left,right)
self.height = 0
def left_height(self):
return self.left.height if self.left is not None else 0
def right_height(self):
return self.right.height if self.right is not None else 0
接下来定义4中调整的非公开方法
def _left_left(self,p):
this = p.node # 有变化的就4个节点
left = this.left
parent = this.parent
left_right = this.left.right
if parent is not None:
if this is parent.left:
parent.left = left
else:
parent.right = left
else:
self._root = left
this.parent = left
left.parent = parent
this.left = left_right
left.right = this
if left_right is not None:
left_right.parent = this
def _right_right(self,p):
this = p.node # 有变化的就4个节点
right = this.right
parent = this.parent
right_left = this.right.left
if parent is not None:
if this is parent.left:
parent.left = right
else:
parent.right = right
else:
self._root = right
this.parent = right
right.parent = parent
this.right = right_left
right.left = this
if right_left is not None:
right_left.parent = this
def _left_right(self,p):
self._right_right(self.left(p))
self._left_left(p)
def _right_left(self,p):
self._left_left(self.right(p))
self._right_right(p)
然后是用于平衡二叉树的方法，也就是根据情况调用上边那4种策略
def _isbalanced(self,p):
"""判断节点是否平衡"""
return abs(p.node.left_height() - p.node.right_height()) <= 1
def _recompute_height(self,p):
"""重新计算高度"""
p.node.height = 1 + max(p.node.left_height(),p.node.right_height())
def _rebalanced(self,p):
while p is not None:
if self._isbalanced(p):
self._recompute_height(p)
p = self.parent(p)
else:
if p.node.left_height()>p.node.right_height() and p.node.left.left_height()>p.node.left.right_height():
# LL的情况，只有自己和左孩子的高度可能变化
self._left_left(p)
elif p.node.right_height()>p.node.left_height() and p.node.right.right_height()>p.node.right.left_height():
# RR的情况，只有自己和右孩子的高度可能变化
self._right_right(p)
elif p.node.left_height()>p.node.right_height() and p.node.left.left_height()<p.node.left.right_height():
# LR的情况，只有自己和左孩子和左孩子的右孩子的高度可能变化
left = self.left(p)
self._left_right(p)
self._recompute_height(left)
else:
# RL的情况，只有自己和右孩子和右孩子的左孩子的高度可能变化
right = self.right(p)
self._right_left(p)
self._recompute_height(right)
while p is not None:
# 调整所有p的祖先的高度
self._recompute_height(p)
p = self.parent(p)
然后把方法封装成删除时和插入时的两个方法，虽然执行的内容是相同的
def _rebalanced_insert(self,p):
"""插入时的平衡调整"""
self._rebalanced(p)
def _rebalanced_delete(self, p):
"""删除时的平衡调整"""
self._rebalanced(p)
最后重写一下setitem方法与删除时调用的方法