486.-merge-k-sorted-arrays.md

486. Merge K Sorted Arrays

class Solution:
    """
    @param arrays: k sorted integer arrays
    @return: a sorted array
    """
    def mergekSortedArrays(self, arrays):
        import heapq
        # write your code here
        heads = [(v[0],i,0)for i,v in enumerate(arrays) if v ]
        heapq.heapify(heads)
        ans = []
        while heads:
            val, idx, i = heapq.heappop(heads)
            # print( val, idx, i)
            ans.append(val)
            if i + 1 == len(arrays[idx]):
                continue
            
            heapq.heappush(heads,(arrays[idx][i+1],idx, i + 1))
            
            
        return ans