heapq 看起来比 PriorityQueue 要快得多,但是用起来也要麻烦一些,尤其是 heapq 是函数式风格,和对象格格不入。
其实可以直接使用 heapq,只要把链表中所有结点的值存进去就行,但如果只想同时维护 top K 个值在里面呢?就需要做一些调整。
直接上代码:
class Solution:
def __init__(self):
setattr(ListNode, '__lt__', lambda self, other: self.val < other.val)
def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
res_head = res = ListNode()
lists_head = [l for l in lists if l]
heapq.heapify(lists_head)
while lists_head:
head = heapq.heappop(lists_head)
res.next = ListNode(head.val)
res = res.next
if head.next:
heapq.heappush(lists_head, head.next)
return res_head.next
相当于是手动指定一个 compartor。为什么可以这样修改,可以看看 heapq 的源码。
def heapify(x):
"""Transform list into a heap, in-place, in O(len(x)) time."""
n = len(x)
# Transform bottom-up. The largest index there's any point to looking at
# is the largest with a child index in-range, so must have 2*i + 1 < n,
# or i < (n-1)/2. If n is even = 2*j, this is (2*j-1)/2 = j-1/2 so
# j-1 is the largest, which is n//2 - 1. If n is odd = 2*j+1, this is
# (2*j+1-1)/2 = j so j-1 is the largest, and that's again n//2-1.
for i in reversed(range(n//2)):
_siftup(x, i)
def _siftup(heap, pos):
endpos = len(heap)
startpos = pos
newitem = heap[pos]
# Bubble up the smaller child until hitting a leaf.
childpos = 2*pos + 1 # leftmost child position
while childpos < endpos:
# Set childpos to index of smaller child.
rightpos = childpos + 1
if rightpos < endpos and not heap[childpos] < heap[rightpos]:
childpos = rightpos
# Move the smaller child up.
heap[pos] = heap[childpos]
pos = childpos
childpos = 2*pos + 1
# The leaf at pos is empty now. Put newitem there, and bubble it up
# to its final resting place (by sifting its parents down).
heap[pos] = newitem
_siftdown(heap, startpos, pos)
这里随便举个建堆的例子。建堆的时候,调用了 heap[childpos] < heap[rightpos]
,如果保存的是对象,会触发对象的比较大小的魔术方法,因此只需要重定义一下类似于 __lt__
这种方法即可。作为脚本语言,python 在修改定义上还是很容易的。
不过,为什么没有提供自定义 comparator 的方法?这才是最奇怪的地方。