原題
將 k 個(gè)排序數(shù)組合并為一個(gè)大的排序數(shù)組。
樣例
給出下面的 3 個(gè)排序數(shù)組：
[
[1, 3, 5, 7],
[2, 4, 6],
[0, 8, 9, 10, 11]
]
合并后的大數(shù)組應(yīng)為：
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
解題思路
- 本題類似于Merge k Sorted Lists。與之不同的是，arrays不能直接通過一個(gè)值找到相鄰的下一個(gè)值，所以需要建立一個(gè)set(數(shù)值，數(shù)組編號(hào)，index)
- 這樣，我們就擁有了當(dāng)前數(shù)組是Arrays中的哪一個(gè)數(shù)組，并且知道了是這個(gè)數(shù)組的第幾個(gè)數(shù)字
- 首先創(chuàng)建一個(gè)最小堆，把每一個(gè)array的第一個(gè)數(shù)字放入堆中，然后取最小，把最小值放入res中，并把最小值所在數(shù)組中的相鄰的數(shù)放入堆中
注意
使用heapq性能優(yōu)于封裝的Queue.PriorityQueue，本題使用Queue.PriorityQueue會(huì)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
