class Solution:
def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
# ListNode.__lt__ = lambda self, other: self.val < other.val # 最简洁的做法是增加属性, 可以把node直接push进heap
heap = []
for i, l in enumerate(lists):
# 注意: 每个list只有head进入heap
if l:
heapq.heappush(heap, (l.val, i, l)) # python3 heap中前一个相等就会判断下一个, 由于node.val可能相等, 需要3个元素或index
dummy = head = ListNode()
while heap:
val, i, node = heapq.heappop(heap)
head.next = node
head = head.next
if node.next:
heapq.heappush(heap, (node.next.val, i, node.next))
return dummy.next
时间复杂度：O(nklog k)
空间复杂度：O(k)
follow up
# 堆一直保存k个数(每个list仅保留一个), 最小的pop后, 相应的list往右移一位, 记录移动过程中heap中最小的范围
class Solution:
def smallestRange(self, nums: List[List[int]]) -> List[int]:
heap = [(row[0], i, 0) for i, row in enumerate(nums)]
heapq.heapify(heap)
max_value = max([row[0] for row in heap])
min_value = min([row[0] for row in heap])
ans = [min_value, max_value]
while len(heap) == len(nums):
num, r, c = heapq.heappop(heap)
if c + 1 < len(nums[r]):
heapq.heappush(heap, (nums[r][c + 1], r, c + 1))
max_value = max(max_value, nums[r][c + 1])
min_value = heapq.nsmallest(1, heap)[0][0]
if max_value - min_value < ans[1] - ans[0]:
ans = [min_value, max_value]
return ans