用python实现堆排序、优先级队列

堆排序是一个优秀的算法，但在实际应用中，性能比快速排序差一些。尽管如此，堆这一数据结构仍然有很多应用。接下来我们要介绍堆的一个常见的应用：优先级队列。

一、堆排序

堆排序图解

数组$A[0,...,A.length-1]$，给定一个节点的下标i，我们很容易计算得到它的父节点、左孩子和右孩子的下标。

def left(i): return 2 * i + 1
def right(i): return 2 * i + 2
def parent(i): return (i-1)//2

以最大堆为例，首先来看最大堆的性质：

$$A[parent(i)]>=A[i]$$
意思就是任意节点的父节点的值必须大于等于该节点自己的值。

维护最大堆的性质：

def max_heapify(A, i, heap_size):
    """
    维护最大堆性质
    :param i: 下标
    :return:
    """

    l = left(i)
    r = right(i)

    if l < heap_size and A[l] > A[i]:
        largest = l
    else:
        largest = i

    if r < heap_size and A[r] > A[largest]:
        largest = r

    if largest != i:
        exchange(i, largest)
        max_heapify(A, largest)

上面函数的输入为一个数组$A$和它的一个下标$i$，以及堆的大小$hea{p}_{s}ize$。在调用$max\_heapify$时，我们假定节点为$left(i)$和$right(i)$的二叉树都是最大堆，但这时$A[i]$有可能小于其孩子，这样就违背了最大堆的性质。$max\_heapify$通过让$A[i]$的值在最大堆中“逐级下降”，从而使得下标$i$为根节点的子树重新遵循最大堆的性质。

建堆：

def build_max_heap(A):
	heap_size = len(A)
    for i in range(heap_size //2, -1, -1):
        max_heapify(A,i,heap_size)

用自底向上的方法，对每个非叶节点（i <= heap_size//2）利用过程$max\_heapify$把一个数组转换为最大堆（本质上还是数组，只不过值的位置发生了变化）。建堆的时间复杂度是$O(n)$。

堆排序算法：

 def heap_sort(A):
     build_max_heap(A)
     heap_size = len(A)
     for i in range(len(A)-1, 0, -1):
         exchange(0, i)
         heap_size -= 1
         max_heapify(A,0,heap_size)

初始时候，建立一个最大堆。然后将第一个元素和最后一个元素交换后，去掉最后一个元素（$heap\_size-=1$)，并使用$max\_heapify(A,0,heap\_size)$保持最大堆性质。不断重复这个过程，直到堆的大小降为1。

至此堆排序算法讲完了。建立一个Heap类，整合上述方法：

class Heap:

    def __init__(self, values):
        if type(values) is not list:
            raise ValueError("{} initial values must be list".format(self.__class__.__name__))
        self.values = values
        self.heap_size = len(values)

    @staticmethod
    def left(i): return 2 * i + 1

    @staticmethod
    def right(i): return 2 * i + 2

    @staticmethod
    def parent(i): return (i-1)//2

    def exchange(self, a, b):
        tmp = self.values[a]
        self.values[a] = self.values[b]
        self.values[b] = tmp

    def max_heapify(self, i):
        """
        维护最大堆性质
        :param i: 下标
        :return:
        """

        l = self.left(i)
        r = self.right(i)

        if l < self.heap_size and self.values[l] > self.values[i]:
            largest = l
        else:
            largest = i

        if r < self.heap_size and self.values[r] > self.values[largest]:
            largest = r

        if largest != i:
            self.exchange(i, largest)
            self.max_heapify(largest)

    def build_max_heap(self):
        for i in range(self.heap_size//2, -1, -1):
            self.max_heapify(i)

    def heap_sort(self):
        self.build_max_heap()
        for i in range(self.heap_size-1, 0, -1):
            self.exchange(0, i)
            self.heap_size -= 1
            self.max_heapify(0)

二、优先级队列

一个最大优先队列Q支持以下操作：

$insert(x)$：把元素$x$插入$Q$队列。
$maximum()$：返回$Q$队列中最大元素。
$pop()$：去掉并返回$Q$队列中最大元素。
$increase\_key(i,x)$：将下标$i$的值增大到$x$。

优先级队列本质也是一个堆，在堆的基础上增加上述方法。

class PriorityQueue(Heap):

    def __init__(self, values):
        super(PriorityQueue, self).__init__(values)
        self.build_max_heap()

    def add(self, value):
        self.values.append(float("-inf"))
        self.heap_size += 1
        self.increase_value(self.heap_size-1, value)

    def pop(self):
        if self.heap_size < 1:
            raise Exception("heap underflow")
        self.exchange(0, self.heap_size-1)
        value = self.values.pop()
        self.heap_size -= 1
        self.max_heapify(0)
        return value

    def maximum(self):
        return self.values[0]

    def increase_value(self, i, value):
        """
        增加第i个元素的值
        :param i: 下标
        :param value:值增加为value
        :return:
        """
        if value < self.values[i]:
            raise Exception("new value is smaller than current value")
        self.values[i] = value
        while i > 0 and self.values[i] > self.values[self.parent(i)]:
            self.exchange(i, self.parent(i))
            i = self.parent(i)

过程$pop$与过程$heap\_sort$ $for$循环体的部分很相似，弹出第一个元素，$heap\_size-=1$，维护堆性质。

值得一提的是过程$increase\_key$，当前元素增大后，不断与其父节点进行比较，如果比父元素值大，则交换它们，直到当前元素小于父元素或没有父元素为止。

而过程$insert$就是在堆底加入了一个负无穷的元素，然后将调用过程$increase\_key$将它增加为输入值的大小。

测试

heap = PriorityQueue([6, 10, 8, 2, 4])
print("intial", ":", heap.values)
heap.add(16)
print("add", 16, ":", heap.values)
heap.add(11)
print("add", 11, ":", heap.values)
heap.pop()
print("pop", ":", heap.values)
print("maximum", ":", heap.maximum())

输出