python实现堆以及堆排序

最新推荐文章于 2024-06-26 19:29:00 发布

NangoGreen

最新推荐文章于 2024-06-26 19:29:00 发布

阅读量1.1k

点赞数 2

分类专栏：数据结构与算法之美（python实现）文章标签： python 数据结构

本文链接：https://blog.csdn.net/NangoGreen/article/details/130995998

版权

数据结构与算法之美（python实现）专栏收录该内容

2 篇文章 0 订阅

订阅专栏

堆作为一种数据结构，它是一种特殊的树。在学习堆之前，我们已经知道了二叉树的概念，而堆就是一种特殊的二叉树，特殊有两点：1.堆必须是一个完全二叉树，也就是说，叶子结点都在最底下两层，最后一层的叶子结点靠左排列，并且除了最后一层，其它层的节点个数要达到最大；2.堆中每一个节点的值都必须大于等于（或小于等于）其子树中每个节点的值，分别称做大顶堆和小顶堆。如下图所示，其中1、2为大顶堆，3为小顶堆，4不是堆。

堆这种数据结构的应用场景非常多，其中最典型的便是堆排序，堆排序是一种原地的，时间复杂度为O(nlogn)的排序算法，在某些场景非常好用。下面就具体讲一下堆以及堆排序的实现。

我们知道，完全二叉树可以用数组来存储，用数组存储非常的节省空间，而且访问数据特别的方便，既然堆是一个完全二叉树，我们就用数组来实现一个堆。对于一个堆，有两个重要操作，一个是在堆中插入数据，一个是删除堆顶数据。

1.在堆中插入一个数据：插入一个数据之后，我们需要使其继续满足堆的特性，我们把插入的数据放在最末尾，这时候它可能就不满足堆的特性了，这时候我们就需要经过一系列比较交换使其重新满足堆的特性，这个过程就叫做堆化，如下图所示。不断找到当前节点的父节点，比较大小判断是否需要进行交换。

2.删除堆顶数据：删除堆顶数据有两种方法，第一种是直接删除堆顶数据，然后找到左右子节点的最大数据顶替堆顶，如此迭代，直到遇到叶子结点。但是这种方式可能会导致最后出现数组空洞从而不满足堆的两个特性，如下图所示。

第二种方法就是先把堆顶数据和最后一个节点交换数据，然后删除最后一个节点，在从上往下进行堆化，这样就不会出现数组空洞了，如下图所示。

下面就讲解一下具体实现代码。这里我们用到了两个模块，一个是random，用于将数据进行打乱（让数据更符合真实情况），另一个是math，在绘制堆的时候需要使用。

首先我们需要创建一个堆类，并且进行初始化，在初始化中，我们要定义一个data数组（列表），用于存储数据，同时定义堆容量。nums作为实例化的时候传入的数组，也就是需要建堆的数据，我们先把它挨个存入data中，然后再进行堆化。这里我们需要定义一个length变量存储数组的总长度，后面需要使用。

初始化之后，我们需要定义一个堆化函数，堆data数组进行堆化。也就是_heapity函数，一个下划线'_'表示这是一个私有方法，不希望被外界访问。然后进行判断，如果数组长度小于等于1的话，说明不需要进行堆化，直接返回。否则，我们让idx=(self._length-2) // 2，idx就记录了堆中最后一个节点父节点的下标。然后我们用一个for循环，从这个节点依次向上进行堆化。例如：假设data中有10个数据，那么idx就等于4，从4到0进行遍历，就可以访问整个堆的所有数据。

接下来是自下向上堆化的具体实现步骤，这里定义一个_heap_down抽象方法，之所以定义抽象方法，是因为在这个代码中，我同时进行了大顶堆和小顶堆的实现，而大顶堆和小顶堆的堆化过程是有所不同的（后面会具体介绍），所以后面需要再创建一个大顶堆和小顶堆的类，让它们继承自Heap类，在子类中去具体实现堆化操作。这里需要注意一点，这个类有抽象方法，所以它是一个抽象类，抽象类中的抽象方法并不做具体实现，而在它的子类中去实现这个抽象方法，子类要进行实例化操作，必须实现父类中的所有抽象方法。如果读者关于这一块不是很了解，建议再补一下python类和继承的相关知识。代码如下

import math
import random


class Heap(object):
    def __init__(self, nums, capacity=100):
        self._data = []
        self._capacity = capacity
        if type(nums) == list and len(nums) <= capacity:
            for n in nums:
                assert type(n) is int
                self._data.append(n)
        self._length = len(self._data)

        self._heapity()  # 堆化

    def _heapity(self):  # 堆化
        if self._length <= 1:  # 如果长度小于等于1，直接返回
            return
        else:
            idx = (self._length - 2) // 2
            for i in range(idx, -1, -1):  # 自下向上堆化
                self._heap_down(i)

    def _heap_down(self, idx):  # 抽象类的抽象方法，在抽象类中并不做实现
        pass

现在我们就来具体实现自下向上堆化，首先我们创建一个大顶堆（MaxHeap）类，继承自Heap类，然后重写父类的_heap_down方法。其方法名与参数必须与父类中一致，让lp等于最后一个节点的父节点，只要idx小于等于lp，就判断其是否需要与左右子节点进行交换，lc表示当前节点的左子节点，rc表示当前节点的右子节点，用tem记录左右子节点较大那一个的索引，然后判断其是否需要与父节点进行交换。交换之后让idx等于被交换节点的索引，因为你把一个数据交换过来之后，还需要判断交换过来的数据与它的子节点的大小，这样才能保证堆的两个特性一直被满足。例如你把索引为0和1的节点交换了数据，还需要判断1和3、4位置数据的大小。如果不需要交换就break，这里一定要有，不然while会进入死循环。

class MaxHeap(Heap):  # 抽象类的子类要想进行实例化，必须实现抽象父类中的所有抽象方法
    def _heap_down(self, idx):  # TODO：（关键）堆化是对整个数组进行
        lp = (self._length - 2) // 2  # 让lp始终指向最后一个节点的父节点
        while idx <= lp:  # 只要当前节点比最后节点的父节点索引小，就得判断是否要进行交换
            lc = 2 * idx + 1  # 当前节点的左节点
            rc = lc + 1  # 当前节点的右节点
            if rc <= self._length - 1:  # 判断当前节点是否有右节点
                tem = lc if self._data[lc] >= self._data[rc] else rc
            else:
                tem = lc
            if self._data[idx] < self._data[tem]:
                self._data[idx], self._data[tem] = self._data[tem], self._data[idx]
                idx = tem  # 防止下一层还需要进行交换
            else:
                break  # 不交换就break，不写这一步会陷入死循环

小顶堆的堆化操作和大顶堆非常类似，只要把两个地方的判断符号改过来即可。代码如下：

class MinHeap(Heap):
    def _heap_down(self, idx):
        lp = (self._length - 2) // 2
        while idx <= lp:
            lc = 2 * idx + 1
            rc = lc + 1
            if rc <= self._length - 1:
                tem = lc if self._data[lc] < self._data[rc] else rc  # tip:小顶堆只需要把这两处的判断符号改过来即可
            else:
                tem = lc
            if self._data[idx] >= self._data[tem]:  # tip:小顶堆只需要把这两处的判断符号改过来即可
                self._data[idx], self._data[tem] = self._data[tem], self._data[idx]
                idx = tem
            else:
                break

堆化完成之后，就开始实现堆的插入数据操作。同样的，大顶堆和小顶堆的插入操作略有不同，所以在父类Heap中写一个insert的抽象方法，在子类中进行实现。其过程与堆化很类似，所不同的是，堆化是对整个堆的所有数据进行了遍历比较，而插入数据只需要让其与父节点，祖父节点等进行比较交换即可，因为插入数据之前，堆以及满足特性了。

    def insert(self, num):  # TODO：（关键）插入只需要与父节点进行比较就可以了
        if self._length >= self._capacity:
            return False
        self._data.append(num)
        lp = len(self._data) - 1
        while lp > 0:
            father = (lp - 1) // 2  # 这一行放在循环里面，就可以保证一直比较到根节点
            if self._data[lp] >= self._data[father]:
                self._data[lp], self._data[father] = self._data[father], self._data[lp]
                lp = father
            else:
                break

小顶堆的插入同样，只需要改响应地方的判断符号即可。

    def insert(self, num):  # TODO：（关键）插入只需要与父节点进行比较就可以了
        if self._length >= self._capacity:
            return False
        self._data.append(num)
        lp = len(self._data) - 1
        while lp > 0:
            father = (lp - 1) // 2
            if self._data[lp] < self._data[father]:  # 判断符号
                self._data[lp], self._data[father] = self._data[father], self._data[lp]
                lp = father
            else:
                break

堆化和插入数据都实现完成了，开始实现删除堆顶数据。这个操作不管是大顶堆还是小顶堆都是一样的，所以我们直接在父类中进行实现，和之前讲的一下，先把堆顶数据和最后节点数据进行交换，然后删除最后一个节点，再调用堆化函数从堆顶开始向下堆化即可。因为堆化操作要使用长度变量，所以这里没删除一次，长度必须减1。

    def _rmv_top(self):
        if len(self._data) == 0:
            return None

        self._data[0], self._data[-1] = self._data[-1], self._data[0]  # 先把根节点和最后一个节点交换
        ret = self._data.pop()
        self._length -= 1  # 长度必须减1
        self._heap_down(0)  # 再从跟节点开始堆化
        return ret

以上就是对一个堆的基本操作的实现过程，做完之后，我们再写一个绘制堆的函数。我们知道，在用print输出一个实例时，会自动调用类中的__repr__内置函数，要绘制堆，我们就对这个内置函数进行重写。这样就可以打印输出了，代码如下。

    def __repr__(self):
        return self._draw_heap(self._data)  # 这里必须要进行返回

    def _draw_heap(self, data):
        if len(data) == 0:
            return 'empty heap'
        ret = ''
        for i, n in enumerate(data):  # enumerate函数可以同时获得索引和值
            ret += str(n)  # 这里要进行强制类型转换，+=相当于append
            if i == 2 ** int(math.log(i + 1, 2) + 1) - 2 or i == len(data) - 1:  # 判断是否要进行换行
                ret += '\n'
            else:
                ret += ', '  # 不换行就加逗号
        return ret  # 返回的是字符串

最后添加程序入口进行测试。

if __name__ == '__main__':
    nums = list(range(10))
    random.shuffle(nums)
    max_heap = MaxHeap(nums)

    print("---max heap---")
    print(max_heap)

    min_heap = MinHeap(nums)
    print("---min heap---")
    print(min_heap)

    max_heap._rmv_top()  # 删除根节点
    print(max_heap)

    min_heap._rmv_top()
    print(min_heap)

    max_heap.insert(10)
    print(max_heap)

    min_heap.insert(6)
    print(min_heap)

整个完整代码如下。

import math
import random


class Heap(object):
    def __init__(self, nums, capacity=100):
        self._data = []
        self._capacity = capacity
        if type(nums) == list and len(nums) <= capacity:
            for n in nums:
                assert type(n) is int
                self._data.append(n)
        self._length = len(self._data)

        self._heapity()  # 堆化

    def _heapity(self):  # 堆化
        if self._length <= 1:  # 如果长度小于等于1，直接返回
            return
        else:
            idx = (self._length - 2) // 2
            for i in range(idx, -1, -1):  # 自下向上堆化
                self._heap_down(i)

    def _heap_down(self, idx):  # 抽象类的抽象方法，在抽象类中并不做实现
        pass

    def _draw_heap(self, data):
        if len(data) == 0:
            return 'empty heap'
        ret = ''
        for i, n in enumerate(data):  # enumerate函数可以同时获得索引和值
            ret += str(n)  # 这里要进行强制类型转换，+=相当于append
            if i == 2 ** int(math.log(i + 1, 2) + 1) - 2 or i == len(data) - 1:  # 判断是否要进行换行
                ret += '\n'
            else:
                ret += ', '  # 不换行就加逗号
        return ret  # 返回的是字符串

    def _rmv_top(self):
        if len(self._data) == 0:
            return None

        self._data[0], self._data[-1] = self._data[-1], self._data[0]  # 先把根节点和最后一个节点交换
        ret = self._data.pop()
        self._length -= 1  # 长度必须减1
        self._heap_down(0)  # 再从跟节点开始堆化
        return ret

    def insert(self, num):  # 插入的抽象方法
        pass

    def __repr__(self):
        return self._draw_heap(self._data)  # 这里必须要进行返回


class MaxHeap(Heap):  # 抽象类的子类要想进行实例化，必须实现抽象父类中的所有抽象方法
    def _heap_down(self, idx):  # TODO：（关键）堆化是对整个数组进行
        lp = (self._length - 2) // 2  # 让lp始终指向最后一个节点的父节点
        while idx <= lp:  # 只要当前节点比最后节点的父节点索引小，就得判断是否要进行交换
            lc = 2 * idx + 1  # 当前节点的左节点
            rc = lc + 1  # 当前节点的右节点
            if rc <= self._length - 1:  # 判断当前节点是否有右节点
                tem = lc if self._data[lc] >= self._data[rc] else rc
            else:
                tem = lc
            if self._data[idx] < self._data[tem]:
                self._data[idx], self._data[tem] = self._data[tem], self._data[idx]
                idx = tem  # 防止下一层还需要进行交换
            else:
                break  # 不交换就break，不写这一步会陷入死循环

    def insert(self, num):  # TODO：（关键）插入只需要与父节点进行比较就可以了
        if self._length >= self._capacity:
            return False
        self._data.append(num)
        lp = len(self._data) - 1
        while lp > 0:
            father = (lp - 1) // 2  # 这一行放在循环里面，就可以保证一直比较到根节点
            if self._data[lp] >= self._data[father]:
                self._data[lp], self._data[father] = self._data[father], self._data[lp]
                lp = father
            else:
                break


class MinHeap(Heap):
    def _heap_down(self, idx):
        lp = (self._length - 2) // 2  # 让lp始终指向最后一个节点的父节点
        while idx <= lp:  # 只要当前节点比最后节点的父节点索引小，就得判断是否要进行交换
            lc = 2 * idx + 1  # 当前节点的左节点
            rc = lc + 1  # 当前节点的右节点
            if rc <= self._length - 1:  # 判断当前节点是否有右节点
                tem = lc if self._data[lc] < self._data[rc] else rc  # tip:小顶堆只需要把这两处的判断符号改过来即可
            else:
                tem = lc
            if self._data[idx] >= self._data[tem]:  # tip:小顶堆只需要把这两处的判断符号改过来即可
                self._data[idx], self._data[tem] = self._data[tem], self._data[idx]
                idx = tem  # 防止下一层还需要进行交换
            else:
                break  # 不交换就break，不写这一步会陷入死循环

    def insert(self, num):  # TODO：（关键）插入只需要与父节点进行比较就可以了
        if self._length >= self._capacity:
            return False
        self._data.append(num)
        lp = len(self._data) - 1
        while lp > 0:
            father = (lp - 1) // 2  # 这一行放在循环里面，就可以保证一直比较到根节点
            if self._data[lp] < self._data[father]:
                self._data[lp], self._data[father] = self._data[father], self._data[lp]
                lp = father
            else:
                break


if __name__ == '__main__':
    nums = list(range(10))
    random.shuffle(nums)
    max_heap = MaxHeap(nums)

    print("---max heap---")
    print(max_heap)

    min_heap = MinHeap(nums)
    print("---min heap---")
    print(min_heap)

    max_heap._rmv_top()  # 删除根节点
    print(max_heap)

    min_heap._rmv_top()
    print(min_heap)

    max_heap.insert(10)
    print(max_heap)

    min_heap.insert(6)
    print(min_heap)

堆排序：堆排序可以分解为两个过程->建堆和排序，建堆我们前面已经讲过了，就是之前的堆化方法，接下来就重点讲一下排序过程。

建堆结束之后，数组中的数据已经是按照大顶堆的特性来组织的。数组中的第一个元素就是堆顶，也就是最大的元素。我们把它跟最后一个元素交换，那最大元素就放到了下标为 n 的位置，然后我们把这个元素删除，再从堆顶堆n-1个元素进行堆化，那么堆顶就变成了n-1个元素中的最大元素。整个过程就是之前所说的删除堆顶元素的过程，不断重复这个过程，直到堆中只有一个元素，这样所得到的数组就按照从小到大排列了，实现代码也非常简单，我们需要把之前写的堆的python文件导入过来，再写一个排序函数即可。

import heap


class SortHeap(heap.MaxHeap):
    def sort(self, nums):
        if len(nums) <= 1:
            return
        for i in range(self._length - 1, -1, -1):
            nums[i] = self._rmv_top()
        return


if __name__ == '__main__':
    nums = [3, 5, 2, 6, 1, 7, 6]
    bhs = SortHeap(nums)

    print('--- before sort ---')
    print(nums)
    print(bhs)

    bhs.sort(nums)
    print('--- after sort ---')
    print(nums)