LintCode 486 [Merge k Sorted Arrays]

原题

将 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