原题
将 k 个排序数组合并为一个大的排序数组。
样例
给出下面的 3 个排序数组：
[
[1, 3, 5, 7],
[2, 4, 6],
[0, 8, 9, 10, 11]
]
合并后的大数组应为：
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
解题思路
- 本题类似于Merge k Sorted Lists。与之不同的是，arrays不能直接通过一个值找到相邻的下一个值，所以需要建立一个set(数值，数组编号，index)
- 这样，我们就拥有了当前数组是Arrays中的哪一个数组，并且知道了是这个数组的第几个数字
- 首先创建一个最小堆，把每一个array的第一个数字放入堆中，然后取最小，把最小值放入res中，并把最小值所在数组中的相邻的数放入堆中
注意
使用heapq性能优于封装的Queue.PriorityQueue，本题使用Queue.PriorityQueue会TLE
完整代码
import heapq
class Solution:
# @param {int[][]} arrays k sorted integer arrays
# @return {int[]} a sorted array
def mergekSortedArrays(self, arrays):
# Write your code here
result = []
heap = []
for index, array in enumerate(arrays):
if len(array) == 0:
continue
heapq.heappush(heap, (array[0], index, 0))
while len(heap):
val, arrayNum, index = heapq.heappop(heap)
result.append(val)
if index + 1 < len(arrays[arrayNum]):
heapq.heappush(heap, (arrays[arrayNum][index + 1], arrayNum, index + 1))
return result
