vlambda博客
学习文章列表

Treap——堆和二叉树的完美结合,性价比极值的搜索树




大家好,今天和大家聊一个新的数据结构,叫做Treap。

Treap本质上也是一颗BST(平衡二叉搜索树),和我们之前介绍的SBT是一样的。但是Treap维持平衡的方法和SBT不太一样,有些许区别,相比来说呢,Treap的原理还要再简单一些,所以之前在竞赛当中不允许使用STL的时候,我们通常都会手写一棵Treap来代替。


Treap的基本原理


既然是平衡二叉搜索树,关键点就在于平衡,那么重点自然是如何维护树的平衡。

在Treap当中,维护平衡非常简单,只有一句话,就是通过维护小顶堆的形式来维持树的平衡。Treap也正是因此得名,因为它是Tree和Heap的结合体。

我们来看下Treap当中节点的结构:

class TreapNode(TreeNode):
    """
    TreeNode: The node class of treap tree.
    Paramters: 
        key: The key of node, can be treated as the key of dictionary
        value: The value of node, can be treated as the value of dictionary
        priority: The priority of node, specially for treap structure, describe the priority of the node in the treap. 
        lchild: The left child of node
        rchild: The right child of node
        father: The parent of node, incase that we need to remove or rotate the node in the treap, so we need father parameter to mark the address of the parent
    """

    def __init__(self, key=None, value=None, lchild=None, rchild=None, father=None, priority=None):
        super().__init__(key, value, lchild, rchild, father)
        self._priority = priority

    @property
    def priority(self):
        return self._priority

    @priority.setter
    def priority(self, priority):
        self._priority = priority

    def __str__(self):
        return 'key={}, value={}'.format(self.key, self.value)

这里的TreeNode是我抽象出来的树结构通用的Node,当中包含key、value、lchild、rchild和father。TreapNode其实就是在此基础上增加了一个priority属性。

之所以要增加这个priority属性是为了维护它堆的性质,通过维护这个堆的性质来保持树的平衡。具体的操作方法,请往下看。


Treap的增删改查

插入

首先来讲Treap的插入元素的操作,其实插入元素的操作非常简单,就是普通BST插入元素的操作。唯一的问题是如何维持树的平衡。

我们前文说了,我们是通过维持堆的性质来保持平衡的,那么自然又会有一个新的问题。为什么维持堆的性质可以保证平衡呢?

答案很简单,因为我们在插入的时候,需要对每一个插入的Node随机附上一个priority。堆就是用来维护这个priority的,保证树根一定拥有最小的priority。正是由于这个priority是随机的,我们可以保证整棵树蜕化成线性的概率降到无穷低。

当我们插入元素之后发现破坏了堆的性质,那么我们需要通过旋转操作来维护。举个简单的例子,在下图当中,如果B节点的priority比D要小,为了保证堆的性质,需要将B和D进行互换。由于直接互换会破坏BST的性质,所以我们采取旋转的操作。

旋转之后我们发现B和D互换了位置,并且旋转之后的A和E的priority都是大于D的,所以旋转之后我们整棵树依然维持了性质。

右旋的情况也是一样的,其实我们观察一下会发现,要交换左孩子和父亲需要右旋,如果是要交换右孩子和父亲,则需要左旋。

整个插入的操作其实就是基础的BST插入过程,加上旋转的判断。

    def _insert(self, node, father, new_node, left_or_right='left'):
        """
        Inside implement of insert node.
        Implement in recursion.
        Since the parameter passed in Python is reference, so when we add node, we need to assign the node to its father, otherwise the reference will lose outside the function.
        When we add node, we need to compare its key with its father's key to make sure it's the lchild or rchild of its father.
        """

        if node is None:
            if new_node.key < father.key:
                father.lchild = new_node
            else:
                father.rchild = new_node
            new_node.father = father
            return
        if new_node.key < node.key:
            self._insert(node.lchild, node, new_node, 'left')
            # maintain
            if node.lchild.priority < node.priority:
                self.rotate_right(node, father, left_or_right)
        else:
            self._insert(node.rchild, node, new_node, 'right')
            # maintain
            if node.rchild.priority < node.priority:
                self.rotate_left(node, father, left_or_right)

前面的逻辑就是BST的插入,也就是和当前节点比大小,决定插入在左边还是右边。注意一下,这里我们在插入完成之后,增加了maintain的逻辑,其实也就是比较一下,刚刚进行的插入是否破坏了堆的性质。可能有些同学要问我了,这里为什么只maintain了一次?有可能插入的priority非常小,需要一直旋转到树根不是吗?

的确如此,但是不要忘了,我们这里的maintain逻辑并非只调用一次。随着整个递归的回溯,在树上的每一层它其实都会执行一次maintain逻辑。所以是可以保证从插入的地方一直维护到树根的。


查询

查询很简单,不用多说,就是BST的查询操作,没有任何变化。

    def _query(self, node, key, backup=None):
        if node is None:
            return backup
        if key < node.key:
            return self._query(node.lchild, key, backup)
        elif key > node.key:
            return self._query(node.rchild, key, backup)
        return node

    def query(self, key, backup=None):
        """
        Return the result of query a specific node, if not exists return None
        """

        return self._query(self.root, key, backup)

删除

删除的操作稍微麻烦了一些,由于涉及到了优先级的维护,不过逻辑也不难理解,只需要牢记需要保证堆的性质即可。

首先,有两种情况非常简单,一种是要删除的节点是叶子节点,这个都很容易想明白,删除它不会影响任何其他节点,直接删除即可。第二种情况是链节点,也就是说它只有一个孩子,那么删除它也不会引起变化,只需要将它的孩子过继给它的父亲,整个堆和BST的性质也不会受到影响。

对于这两种情况之外,我们就没办法直接删除了,因为必然会影响堆的性质。这里有一个很巧妙的做法,就是可以先将要删除的节点旋转,将它旋转成叶子节点或者是链节点,再进行删除。

在这个过程当中,我们需要比较一下它两个孩子的优先级,确保堆的性质不会受到破坏。

    def _delete_node(self, node, father, key, child='left'):
        """
        Implement function of delete node.
        Defined as a private function that only can be called inside.
        """

        if node is None:
            return
        if key < node.key:
            self._delete_node(node.lchild, node, key)
        elif key > node.key:
            self._delete_node(node.rchild, node, key, 'right')
        else:
            # 如果是链节点,叶子节点的情况也包括了
            if node.lchild is None:
                self.reset_child(father, node.rchild, child)
            elif node.rchild is None:
                self.reset_child(father, node.lchild, child)
            else:
                # 根据两个孩子的priority决定是左旋还是右旋
                if node.lchild.priority < node.rchild.priority:
                    node = self.rotate_right(node, father, child)
                    self._delete_node(node.rchild, node, key, 'right')
                else:
                    node = self.rotate_left(node, father, child)
                    self._delete_node(node.lchild, node, key)

                    
    def delete(self, key):
        """
        Interface of delete method face outside.
        """

        self._delete_node(self.root, None, key, 'left')

修改

修改的操作也非常简单,我们直接查找到对应的节点,修改它的value即可。


旋转

我们也贴一下旋转操作的代码,其实这里的逻辑和之前SBT当中介绍的旋转操作是一样的,代码也基本相同:

    def reset_child(self, node, child, left_or_right='left'):
        """
        Reset the child of father, since in Python all the instances passed by reference, so we need to set the node as a child of its father node.
        """

        if node is None:
            self.root = child
            self.root.father = None
            return
        if left_or_right == 'left':
            node.lchild = child
        else:
            node.rchild = child
        if child is not None:
            child.father = node


 def rotate_left(self, node, father, left_or_right):
        """
        Left rotate operation of Treap.
        Example: 

                D
              /   \
             A      B
                   / \
                  E   C

        After rotate:

                B
               / \
              D   C
             / \
            A   E 
        """

        rchild = node.rchild
        node.rchild = rchild.lchild
        if rchild.lchild is not None:
            rchild.lchild.father = node
        rchild.lchild = node
        node.father = rchild
        self.reset_child(father, rchild, left_or_right)
        return rchild

    def rotate_right(self, node, father, left_or_right):
        """
        Right rotate operation of Treap.
        Example: 

                D
              /   \
             A     B
            / \
           E   C

        After rotate:

                A
               / \
              E   D
                 / \
                C   B 
        """

        lchild = node.lchild
        node.lchild = lchild.rchild
        if lchild.rchild is not None:
            lchild.rchild.father = node
        lchild.rchild = node
        node.father = lchild
        self.reset_child(father, lchild, left_or_right)
        return lchild


后记


基本上到这里整个Treap的原理就介绍完了,当然除了我们刚才介绍的基本操作之外,Treap还有一些其他的操作。比如可以split成两个Treap,也可以由两个Treap合并成一个。还可以查找第K大的元素,等等。这些额外的操作,我用得也不多,就不多介绍了,大家感兴趣可以去了解一下。

Treap这个数据结构在实际当中几乎没有用到过,一般还是以竞赛场景为主,我们学习它主要就是为了提升和锻炼我们的数据结构能力以及代码实现能力。Treap它的最大优点就是实现简单,没有太多复杂的操作,但是我们前面也说了,它是通过随机的priority来控制树的平衡的,那么它显然无法做到完美平衡,只能做到不落入最坏的情况,但是无法保证可以进入最好的情况。不过对于二叉树来说,树深的一点差距相差并不大。所以Treap的性能倒也没有那么差劲,属于一个性价比非常高的数据结构。

最后,还是老规矩,我把完整的代码放在了paste当中,大家感兴趣可以点击阅读原文查看,代码里都有详细的注释,大家应该都能看明白。

今天的文章就到这里,衷心祝愿大家每天都有所收获。如果还喜欢今天的内容的话,请来一个三连支持吧~(点赞、在看、转发)