Python递归实现树堆数据结构

2023-04-16 00:00:00 python 递归数据结构

树堆是一种自平衡的树形数据结构，类似于二叉堆，但是一个节点可以有多个子节点。每个节点有一个父节点指针，以及一个实数键值。堆是一种特殊的树，满足堆的性质，即任何一个节点的键值都大于（或小于）它的子节点的键值，且堆是完全二叉树。在树堆中，节点的左子树包含比节点键值小的所有节点，右子树包含比节点键值大的所有节点。

Python递归实现树堆数据结构的代码如下：

class TreeNode:
    def __init__(self, val, parent=None):
        self.val = val
        self.children = []
        self.parent = parent

    def __str__(self):
        return str(self.val)

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

    def insert(self, val):
        if not self.root:
            self.root = TreeNode(val)
        else:
            self._insert(self.root, val)
            self._heapify_inserted_node(val)

    def _insert(self, node, val):
        if val < node.val:
            node.children.append(TreeNode(node.val, parent=node))
            node.val = val
            return True
        elif len(node.children)==0:
            node.children.append(TreeNode(val, parent=node))
            return True
        else:
            return self._insert(node.children[-1], val)

    def _heapify_inserted_node(self, val):
        node = self._find_node(self.root, val)
        while node.parent and node.val < node.parent.val:
            node.val, node.parent.val = node.parent.val, node.val
            node = node.parent

    def _find_node(self, node, val):
        if node.val == val:
            return node
        elif len(node.children) == 0:
            return None
        else:
            for child in node.children:
                found = self._find_node(child, val)
                if found:
                    return found
            return None

    def remove(self):
        if not self.root:
            return None
        elif len(self.root.children) == 0:
            val = self.root.val
            self.root = None
            return val
        else:
            val = self.root.val
            self.root.val = self._remove(self.root.children[-1])
            self._heapify_removed_node(self.root)
            return val

    def _remove(self, node):
        if len(node.children)==0:
            node.parent.children.remove(node)
            return node.val
        else:
            return self._remove(node.children[-1])

    def _heapify_removed_node(self, node):
        for child in node.children:
            if child.val < node.val:
                node.val, child.val = child.val, node.val
                self._heapify_removed_node(child)

其中，TreeNode表示树堆的节点，包含键值、子节点列表以及父节点指针。TreeHeap表示树堆数据结构，包含一个根节点，并提供插入和删除节点的方法。在插入节点时，首先判断是否存在根节点，若不存在，则将插入节点作为根节点；否则，将插入节点和根节点比较大小，若插入节点的值小于根节点，则将插入节点作为新的根节点，将原先的根节点作为插入节点的子节点；若插入节点比根节点大，则递归判断插入节点应该插入哪个子节点中。在插入节点后，需要将堆恢复性质，即将插入节点不断与其父节点比较，若小于父节点，则交换值，并继续与祖先节点比较，直到找到一个大于等于该节点的祖先节点或已到达根节点。

在删除节点时，需要将根节点的值保存下来，并将根节点的值替换为最后一个子节点的值，然后将该子节点删除。同时，需要将根节点不断与其子节点比较，若大于子节点，则交换值，并继续与子节点比较，直到找到一个小于等于该节点的子节点或已到达叶子节点。

树堆的时间复杂度为O(log n)，比较常用的应用是求第k大的元素。

相关文章