Python中的树形数据结构: 树堆

2023-04-11 00:00:00 python 数据结构

树堆是一种基于堆的树形数据结构，可以在O(log n)的时间内完成插入、删除、查找最小值等操作。

在树堆中，每个节点都有一个父节点和若干个子节点。节点之间的关系不仅仅是父子关系，还包括兄弟关系和祖先后代关系。树堆中的每个节点都有一个权值，根据权值大小来比较节点的大小。

根据树堆的定义，我们可以使用以下代码实现一个简单的树堆类：

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.rank = 1
        self.left = None
        self.right = None

class TreeHeap:
    def __init__(self):
        self.root = None

    def __add_node(self, node1, node2):
        if not node1:
            return node2
        if not node2:
            return node1
        if node1.val < node2.val:
            node1.right = self.__add_node(node1.right, node2)
            node1.rank = max(node1.rank, node1.right.rank + 1)
            return node1
        else:
            node2.right = self.__add_node(node2.right, node1)
            node2.rank = max(node2.rank, node2.right.rank + 1)
            return node2

    def push(self, val):
        new_node = TreeNode(val)
        self.root = self.__add_node(self.root, new_node)

    def pop(self):
        if not self.root:
            return None
        min_val = self.root.val
        if self.root.left:
            self.root = self.__add_node(self.root.left, self.root.right)
        else:
            self.root = None
        return min_val

    def get_min(self):
        if not self.root:
            return None
        return self.root.val

以上代码实现了树堆的基本操作，包括插入(push)、删除最小值(pop)、获取最小值(get_min)等。具体实现方式是先定义一个节点类(TreeNode)，其中包含节点的值(val)、子树高度(rank)、左子节点(left)和右子节点(right)等信息；然后定义一个树堆类(TreeHeap)，其中包含根节点(root)和三个基本操作。

下面是一些使用树堆的例子：

heap = TreeHeap()
heap.push(2)
heap.push(1)
heap.push(3)
print(heap.get_min()) # output: 1
print(heap.pop()) # output: 1
print(heap.get_min()) # output: 2

以上代码演示了如何使用树堆来排序，先插入三个数2、1、3，然后发现最小值是1，于是弹出1，最后输出新的最小值2。

相关文章