python--lintcode486.合并k个排序数组

描述

将 k 个有序数组合并为一个大的有序数组。

样例

样例 1:

输入:
  [
    [1, 3, 5, 7],
    [2, 4, 6],
    [0, 8, 9, 10, 11]
  ]
输出: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

样例 2:

输入: 
  [
    [1,2,3],
    [1,2]
  ]
输出: [1,1,2,2,3]

挑战

在 O(N log k) 的时间复杂度内完成：

• N 是所有数组包含的整数总数量。
• k 是数组的个数。

这一题的解法就是先把每个数组的第一个数拿出来，构建小根堆，然后每次从小根堆里拿出最小的数，相应数组的指针往后，取出下一个数重新构建小根堆。这样每次构建小根堆的时间复杂度为logk，总共n个数的话，总时间复杂度为nlogk。使用python的heapq来完成，代码如下：

class Solution:
    """
    @param arrays: k sorted integer arrays
    @return: a sorted array
    """
    def mergekSortedArrays(self, arrays):
        # write your code here
        import  heapq
        result = []
        heap = []
        for index, array in enumerate(arrays):
            if len(array) == 0:
                continue
            heapq.heappush(heap, (array[0], index, 0))

        while len(heap):
            val, x, y = heap[0]
            heapq.heappop(heap)
            result.append(val)
            if y + 1 < len(arrays[x]):
                heapq.heappush(heap, (arrays[x][y + 1], x, y + 1))

        return result


arrays=[
            [1, 3, 5, 7],
            [2, 4, 6],
            [0, 8, 9, 10, 11]
       ]
s = Solution()
print(s.mergekSortedArrays(arrays))