python 中的堆 （heapq 模块）应用：Merge K Sorted Lists

堆是计算机科学中一类特殊的数据结构的统称。堆通常是一个可以被看做一棵树的数组对象。在队列中，调度程序反复提取队列中第一个作业并运行，因为实际情况中某些时间较短的任务将等待很长时间才能结束，或者某些不短小，但具有重要性的作业，同样应当具有优先权。堆即为解决此类问题设计的一种数据结构。

1 定义

n个元素序列{k1,k2...ki...kn},当且仅当满足下列关系时称之为堆：
(ki <= k2i, ki <= k2i+1)或者(ki >= k2i, ki >= k2i+1), (i = 1,2,3,4...n/2)

2 性质

堆的实现通过构造二叉堆（binary heap），实为二叉树的一种；由于其应用的普遍性，当不加限定时，均指该数据结构的这种实现。这种数据结构具有以下性质。

• 任意节点小于（或大于）它的所有后裔，最小元（或最大元）在堆的根上（堆序性）。
• 堆总是一棵完全树。即除了最底层，其他层的节点都被元素填满，且最底层尽可能地从左到右填入。

将根节点最大的堆叫做最大堆或大根堆，根节点最小的堆叫做最小堆或小根堆。常见的堆有二叉堆、斐波那契堆等。

3 应用

3.1 heapq 模块中的重要函数

>>> import heapq
>>> heapq.__all__
['heappush', 'heappop', 'heapify', 'heapreplace', 'merge', 'nlargest', 'nsmallest', 'heappushpop']

heappush(heap, x)       将 x 入堆
heappop(heap)           将堆属性中的最小元素弹出   
heapify(heap)           将 heap 属性强制应用到任意一个列表
heapreplace(heap, x)    将堆中的最小的元素弹出，同时将 x 入堆
merge()                 多个堆合并         
nlargest(n, iter)       返回 iter 中的第 n 大的元素
nsmallest(n, iter)      返回 iter 中的第 n 小的元素

3.2 创建堆结构

list1 = [3,2,8,1,4,5]
heapq.heapify(list1)

list2 = []
heapq.heappush(list2, 3)
heapq.heappush(list2, 2)
heapq.heappush(list2, 8)
heapq.heappush(list2, 1)
heapq.heappush(list2, 4)
heapq.heappush(list2, 5)

list1
Out[110]: [1, 2, 5, 3, 4, 8]

list2
Out[108]: [1, 2, 5, 3, 4, 8]

以上两种方法创建的堆是一样的，元素推入堆后会自动按照二叉树的规范重新排序

3.3 Merge K sorted Lists

用堆的思想合并k个排序链表，并且返回合并后的排序链表。

思路1：

将所有链表的节点 push 到堆中，每次把最小的 pop 出来。代码如下：

"""
Definition of ListNode
class ListNode(object):

    def __init__(self, val, next=None):
        self.val = val
        self.next = next
"""

import heapq
class Solution(object):
    def mergeKLists(self, lists):
        """
        :type lists: List[ListNode]
        :rtype: ListNode
        """
        heap = []
        for node in lists:
            while node:
                heapq.heappush(heap, node.val)
                node = node.next
                
        temp = ListNode(-1)
        head = temp
        while heap:
            smallestNode_val = heapq.heappop(heap)
            temp.next = ListNode(smallestNode_val)
            temp = temp.next
        
        return head.next

思路2：

不用堆用列表也能实现，代码如下：

"""
Definition of ListNode
class ListNode(object):

    def __init__(self, val, next=None):
        self.val = val
        self.next = next
"""

class Solution(object):
    def mergeKLists(self, lists):
        """
        :type lists: List[ListNode]
        :rtype: ListNode
        """
        l = []
        for node in lists:
            while node:
                l.append(node.val)
                node = node.next
        l.sort()
        temp = ListNode(-1)
        head = temp
        for i in l:
            temp.next = ListNode(i)
            temp = temp.next

        return head.next

参考资料：

heapq — Heap queue algorithm