伸展树的python实现

我的代码片:https://code.csdn.net/snippets_manage

# 定义顶点
class Vertex:
    def __init__(self, key, sum, left, right, parent):
        (self.key, self.sum, self.left, self.right, self.parent) = (key, sum, left, right, parent)

# 更新结点的sum属性
def update(v):
    if v == None:
        return
    # v.sum = v.key + v.left.sum + v.right.sum
    v.sum = v.key + (v.left.sum if v.left != None else 0) + (v.right.sum if v.right != None else 0)
    if v.left != None:
        v.left.parent = v
    if v.right != None:
        v.right.parent = v

# 左旋和右旋
def smallRotation(v):
    parent = v.parent
    if parent == None:
        return
    grandparent = v.parent.parent
    # 右旋
    if parent.left == v:
        m = v.right
        v.right = parent
        parent.left = m
    # 左旋 
    else:
        m = v.left
        v.left = parent
        parent.right = m
    update(parent)
    update(v)
    v.parent = grandparent 
    if grandparent != None:
        if grandparent.left == parent:
            grandparent.left = v
        else: 
            grandparent.right = v

# zig-zig,zig-zag
def bigRotation(v):
    if v.parent.left == v and v.parent.parent.left == v.parent:
        # Zig-zig
        smallRotation(v.parent)
            smallRotation(v)
    elif v.parent.right == v and v.parent.parent.right == v.parent:
        # Zig-zig
        smallRotation(v.parent)
        smallRotation(v)    
    else: 
        # Zig-zag
        smallRotation(v);
        smallRotation(v);


# 将结点v旋转至根节点        
def splay(v):
    if v == None:
        return None
    while v.parent != None:
        if v.parent.parent == None:
            smallRotation(v)
        break
        bigRotation(v)
    return v

# 搜索伸展树中键值为key的结点
# 若结点存在,返回该结点
# 若该结点不存在,返回大于该结点的最小结点
# 若key对应的结点是树中的最大结点,返回None
def find(root, key):
    v = root
    last = root
    next = None
    while v is not None:
        if v.key >= key and (next is None or v.key < next.key):
            next = v
        last = v
        if v.key == key:
            break
        if v.key < key:
            v = v.right
        else:
            v = v.left
    root = splay(last)
    return (next, root)

# 把伸展树分割成两棵子树left,right
# left子树的结点的键值均小于key,right的结点的键值均大于等于key
def split(root, key):
    (result, root) = find(root, key)
    if result is None:
        return (root, None)
    right = splay(result)
    left = right.left
    right.left = None
    if left is not None:
        left.parent = None
    update(left)
    update(right)
    return (left, right)


# 将两棵子树left和right合成一颗伸展树返回
def merge(left, right):
    if left is None:
        return right
    if right is None:
        return left
    while right.left is not None:
        right = right.left
    right = splay(right)
    right.left = left
    update(right)
    return right



# 往伸展树中插入一个结点
def insert(x):
    global root
    (left, right) = split(root, x)
    new_vertex = None
    if right is None or right.key != x:
        new_vertex = Vertex(x, x, None, None, None)
    root = merge(merge(left, new_vertex), right)
    (N, root) = find(root, x)

# 删除键值为x的结点
def erase(x):
    global root
    # Implement erase yourself
    (left, right) = split(root, x)
    if right is None or right.key != x:
        root = merge(left, right)
        return
    right = right.right
    if right is None:
        root = merge(left, right)
        return
    right.parent = None
    root = merge(left, right)
    return