使用python3中的heapq模块合并k个排序列表

Question

问题：- 合并 k 个排序列表。

我想使用最小堆来解决这个问题，它可以通过Python中的heapq模块来实现。 以下是该函数的示例代码...

heapq.heappush(listwithoutNone,(node.val, node))
        

但问题是 python 解释器引发错误：

类型错误：'<' not supported between instances of 'ListNode' and 'ListNode' heapq.heappush(listwithoutNone,(node.val, node))

所以，我想使用node.val作为minheap节点元素，因为我传递了一个元组，所以我应该在代码中做什么改变，以便minheap将用node.val堆化堆。

提前致谢。

Answer 1

比较元组时，会比较它们的第一个元素，然后使用它们的第二个元素、它们的元素等来打破任何联系。例如，

(2, "a") < (2, "b")

将计算为

True

。

在这里，您将

(node.val, node)

元组插入到堆中，该堆尝试对它们进行比较。如果节点值存在平局，它将移至元组的第二个元素 - 节点本身。这些只是

ListNode

实例。 Python 不知道如何比较两个

ListNodes

，因此会出现错误。

要启用

ListNodes

之间的比较，您需要实现丰富的比较方法。 一种快速的方法是简单地实现

ListNode.__lt__

，然后使用

functools.total_ordering

装饰器：

import heapq
from functools import total_ordering


@total_ordering
class ListNode:
    def __init__(self, value: float, label: str) -> None:
        self.value = value
        self.label = label

    def __lt__(self, other: 'ListNode'):
        return self.value <= other.value

    def __str__(self):
        return f"ListNode(label={self.label}, value={self.value})"



nodes = []
a =  ListNode(5, "A")
b = ListNode(3, "B")
c = ListNode(5, "C")
heapq.heappush(nodes, a)
heapq.heappush(nodes, b)
heapq.heappush(nodes, c)

while nodes:
    x = heapq.heappop(nodes)
    print(x)

这里我们说比较两个

ListNode

对象与比较它们的值是一样的。定义比较后，您甚至根本不需要插入元组。您可以直接插入

ListNode

对象，并依赖比较方法。

Answer 2

我想你正在谈论这个：Leetcode Merge k Sorted Lists

我正在为您分享一个可行的解决方案：

head = curr = ListNode(0)    # Creating a dummy node
    lst = []
    for l in lists:
        if l:
 # Here we need to first push val so that heap can know which is minimum and which is maximum. 
 # If we insert directly node then heap can't work proper way (in this case heapq.heappop doesn't return minimum node).    
            lst.append((l.val,l))    

    
    heapq.heapify(lst)
    while lst:
        val,node = heapq.heappop(lst)
        curr.next = node
        curr = curr.next
        node = node.next
        if node:
            heapq.heappush(lst,(node.val,node))
    return head.next 

Answer 3

为了改进 Saurabh 的答案，您可以在迭代各个列表时生成最终的排序列表。

这个想法是首先使用每个已排序数组的第一个元素填充堆，该数组标记有该元素来自的数组，并且每当从堆中弹出时，添加与刚刚从堆中弹出的元素关联的数组的下一个元素。

def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
    heap = []
    for i, l in enumerate(lists):
        if l is None:
            continue
        heapq.heappush(heap, (l.val,i))
        lists[i] = lists[i].next

    out = ListNode() # dummy
    cur = out
    while heap:
        val, i = heapq.heappop(heap)
        cur.next = ListNode(val)
        cur = cur.next
        if lists[i] is not None:
            heapq.heappush(heap, (lists[i].val, i))
            lists[i] = lists[i].next
    return out.next

这也可以修改，以便堆跟踪实际节点，这样就不会创建新节点（除了虚拟头之外）：

def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
    heap = []
    for i, l in enumerate(lists):
        if l is None:
            continue
        heapq.heappush(heap, (l.val, i, l))
        lists[i] = lists[i].next

    out = ListNode() # dummy
    cur = out
    while heap:
        val, i, node = heapq.heappop(heap)
        cur.next = node
        cur = cur.next
        if lists[i] is not None:
            heapq.heappush(heap, (lists[i].val, i, lists[i]))
            lists[i] = lists[i].next
    return out.next

堆最多只能包含

k

个元素，其中

k

是输入列表的数量。因此，我们解决了 @Abhishek Jaiswal 的担忧，即空间复杂度应该是

O(k)

。

编辑：添加空间复杂度改进分析

Answer 4

您需要重写 ListNode 提供的 lt 函数。 这是代码

def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
    def my_lt(self, other):
        return self.val<other.val
    setattr(ListNode, "__lt__", my_lt)
    heap = []
    for l in lists:
        if l:
            heap.append((l.val, l))
    heapq.heapify(heap)
    head = ListNode(0)
    curr = head

    while heap:
        val, node = heapq.heappop(heap)
        curr.next = node
        curr = curr.next
        if node.next:
            heapq.heappush(heap, (node.next.val, node.next))
    
    return head.next