Python实现 《算法导论 第三版》中的算法 第6章 堆排序

版权声明：本文为博主原创文章，如有转载请标注来源。 https://blog.csdn.net/shengchaohua163/article/details/83038413

第6章 堆排序

1. 堆

堆是一个数组，它可以被看成一个近似的完全二叉树。树上的每一个结点对应数组中的一个元素。除最底层外，该树是完全充满的，而且是从左向右填充。表示堆的数组包括两个属性：A.length(通常)给出数组元素的个数，A.heap-size表示有多少堆元素存储在该数组中。也就是说，虽然A[1…A.length]可能都存有数据，但只有A[1…A.heap-size]中存放的是堆的有效元素。这里0<=A.heap-size<=A.length。

下面实现了书中的伪代码和练习题，包括：

• P84：计算某一结点的父结点，左孩子和右孩子的下标见前三个函数。
• P86：MAX-HEAPIFY，该函数用于维护最大堆的性质，而且需要传入参数heap_size。递归版本和循环版本分别见max_heapify_recursive和max_heapify_loop。
• P87：练习6.2-5，用循环控制结构取代递归重写MAX-HEAPIFY，见max_heapify_loop。
• P87：BUILD-MAX-HEAP，采用自底向上的方法利用过程MAX-HEAPIFY把一个大小为n=A.length的数组A[1…n]转换为最大堆，而且需要传入参数heap_size。
• P89：HEAPSORT，堆排序算法利用BUILD-MAX-HEAP将输入数组A[1…n]建成最大堆，其中n=A.length，见heap_sort。

简单说一下这几个函数的原理：

• MAX-HEAPIFY：如果以某一结点 $i$为根结点的子树不满足最大堆性质，则将其向下筛选（逐级下降）：与较大的子结点进行比较，若小于则进行交换，递归（或循环）判断子结点是否满足最大堆性质，满足则退出，不满足则继续进行筛选，最终以结点 $i$为根结点的子树会满足最大堆性质。时间复杂度为 $O(lgn)$。
• BUILD-MAX-HEAP：采用自底向上的方法利用MAX-HEAPIFY把一个大小为n=A.length的数组A[1…n]转换为最大堆。因为最大堆中的每个叶结点都可以看成只包含一个元素的堆，所以可以从最后一个非叶结点（自底向上）遍历到根结点，对每个结点都调用一次MAX-HEAPIFY。时间复杂度为 $O(n)$。
• HEAPSORT：堆排序算法利用BUILD-MAX-HEAP将输入数组A[1…n]建成最大堆，其中n=A.length。然后重复这一过程：把最大元素与堆中最后一个元素进行交换，然后去掉最后一个结点（通过减小A.heap_size）。由于新的根结点不满足最大堆，调用一次MAX-HEAPIFY得到新的最大堆。直到堆中只剩一个元素位置。时间复杂度为 $O(nlgn)$。

def get_parent(i):
    return (i-1) // 2


def get_left(i):
    return 2 * i + 1


def get_right(i):
    return 2 * i + 2


def max_heapify_recursive(A, heap_size, i):
    l = get_left(i)
    r = get_right(i)
    largest_ind = i
    if l < heap_size and A[l] > A[largest_ind]:
        largest_ind = l
    if r < heap_size and A[r] > A[largest_ind]:
        largest_ind = r
    if largest_ind != i:
        A[i], A[largest_ind] = A[largest_ind], A[i]
        max_heapify_recursive(A, heap_size, largest_ind)
    
    
def max_heapify_loop(A, heap_size, i): 
    while i < heap_size:
        l = get_left(i)
        r = get_right(i)
        largest_ind = i
        if l < heap_size and A[l] > A[largest_ind]:
            largest_ind = l
        if r < heap_size and A[r] > A[largest_ind]:
            largest_ind = r
        if largest_ind == i:
            break
        else:
            A[i], A[largest_ind] = A[largest_ind], A[i]
            i = largest_ind


def build_max_heap(A, heap_size):
    begin = len(A)//2 - 1  # len(A)//2 - 1是堆中第一个叶子节点的前一个节点
    for i in range(begin, -1, -1):
        max_heapify_loop(A, heap_size, i)


def heap_sort(A):
    heap_size = len(A)
    build_max_heap(A, heap_size)
    for i in range(len(A)-1, 0, -1):
        A[0], A[i] = A[i], A[0]
        heap_size -= 1
        max_heapify_loop(A, heap_size, 0)
    

def test():
    A = [16, 4, 10, 14, 7, 9, 3, 2, 8, 1] # P86例子
    max_heapify_recursive(A, len(A), 1)
    print(A)
    
    A = [16, 4, 10, 14, 7, 9, 3, 2, 8, 1] # P86例子
    max_heapify_loop(A, len(A), 1)
    print(A)
    
    B = [4, 1, 3, 2, 16, 9, 10, 14, 8, 7] # P88例子
    build_max_heap(B, len(B))
    print(B)
    
    C = [16, 14, 10, 8, 7, 9, 3, 2, 4, 1] # P89例子
    heap_sort(C)
    print(C)
    
    
if __name__ == '__main__':
    test()

2. 优先队列

堆的一个常见应用是作为高效的优先队列。优先队列也有两种形式：最大优先队列和最小优先队列。这里我们关注如何寄语最大堆实现最大优先队列。优先队列是一种用来维护由一组元素构成的集合S的数据结构，其中每一个元素都有一个相关的值，成为关键字(key)。

下面将最大堆的代码进行了封装，将包含堆元素的数组和堆的大小作为私有成员变量，将上述实现最大堆的函数作为私有成员函数。为了便于观察，添加了两个输出函数，print_max_heap输出堆元素，print_all输出整个数组。然后实现了书中优先队列部分的伪代码，包括：

• P91：HEAP-MAXIMUM，返回堆中最大元素。时间复杂度为 $\Theta(1)$。
• P91：HEAP-EXTRACT-MAX，去除堆中最大元素，同时返回该元素。时间复杂度为 $O(lgn)$。
• P91：HEAP-INCREASE-KEY，将结点 $i$的值更新为新值。时间复杂度为 $O(lgn)$。
• P92：MAX-HEAP-INSERT，插入操作。时间复杂度为 $O(lgn)$。

class MaxHeap:
    def __init__(self, A):
        self.__A = A
        self.__heap_size = len(A) 
        self.__build_max_heap()
        
    def __get_parent(self, i):
        return (i-1) // 2
    
    def __get_left(self, i):
        return 2 * i + 1
    
    def __get_right(self, i):
        return 2 * i + 2
    
    def __max_heapify_recursive(self, i):
        l = self.__get_left(i)
        r = self.__get_right(i)
        largest_ind = i
        if l <= self.__heap_size and self.__A[l] > self.__A[largest_ind]:
            largest_ind = l
        if r <= self.__heap_size and self.__A[r] > self.__A[largest_ind]:
            largest_ind = r
        if largest_ind != i:
            self.__A[i], self.__A[largest_ind] = self.__A[largest_ind], self.__A[i]
            self.__A = self.__max_heapify_recursive(largest_ind)
        
    def __max_heapify_loop(self, i):
        while i < self.__heap_size:
            l = self.__get_left(i)
            r = self.__get_right(i)
            largest_ind = i
            if l < self.__heap_size and self.__A[l] > self.__A[largest_ind]:
                largest_ind = l
            if r < self.__heap_size and self.__A[r] > self.__A[largest_ind]:
                largest_ind = r
            if largest_ind == i:
                break
            else:
                self.__A[i], self.__A[largest_ind] = self.__A[largest_ind], self.__A[i]
                i = largest_ind
    
    def __build_max_heap(self):
        if len(self.__A) == 1:
            return None
        begin = len(self.__A)//2 - 1  # len(A)//2 - 1是堆中第一个叶子节点的前一个节点
        for i in range(begin, -1, -1):
            self.__max_heapify_loop(i)
    
    def __heap_sort(self):
        self.__build_max_heap()
        for i in range(len(self.__A)-1, 0, -1):
            self.__A[0], self.__A[i] = self.__A[i], self.__A[0]
            self.__heap_size -= 1
            self.__max_heapify_loop(0)
            
    def print_max_heap(self):
        print(self.__A[:self.__heap_size])
        
    def print_all(self):
        print(self.__A)
        
    def heap_maximum(self):
        if self.__heap_size == 0:
            print("Error! Heap is empty!")
        
        return self.__A[0]
        
    def heap_extract_max(self):
        if self.__heap_size == 0:
            print("Error! Heap is empty!")
         
        max = self.__A[0]
        self.__A[0] = self.__A[self.__heap_size-1]
        self.__A.pop() # 删除最后一个元素（书上没有）
        self.__heap_size -= 1 
        self.__max_heapify_loop(0)
        return max
     
    def heap_increase_key(self, i, key):
        if i < 0:
            return
        
        if key < self.__A[i]:
            print("New key is smaller than current key.")
            
        self.__A[i] = key
        while i > 0 and self.__A[self.__get_parent(i)] < self.__A[i]:
            self.__A[i], self.__A[self.__get_parent(i)] = \
                self.__A[self.__get_parent(i)], self.__A[i]
            i = self.__get_parent(i)
            
    def max_heap_insert(self, key):
        import math
        self.__A.append(-math.inf)
        self.__heap_size += 1
        self.heap_increase_key(self.__heap_size-1, key)    

def test():
    A = [16, 4, 10, 14, 7, 9, 3, 2, 8, 1] # P86例子
    mh = MaxHeap(A) # 定义一个最大堆，可以实现优先队列的功能
    mh.print_max_heap() # 等价于print(A)
    print(mh.heap_maximum())
    print(mh.heap_extract_max())
    mh.print_max_heap()
    
    A = [16, 4, 10, 14, 7, 9, 3, 2, 8, 1] # P86例子
    mh = MaxHeap(A)
    mh.heap_increase_key(8, 15)
    mh.print_max_heap()
    mh.max_heap_insert(11)
    mh.print_max_heap()
   
if __name__ == '__main__':
    test()