Merge K Sorted Arrays – Python 3 (Week 16 – 03)

import heapq

class Solution:
    """
    @param arrays: k sorted integer arrays
    @return: a sorted array
    """
    def mergekSortedArrays(self, arrays):
        # write your code here
        result = []
        heap = []
        for index, array in enumerate(arrays):
            if len(array) == 0:
                continue
            #(value, index_of_array, index_of_element)
            heapq.heappush(heap, (array[0], index, 0))
            
        while len(heap):
            val, x, y = heapq.heappop(heap)
            result.append(val)
            if y == len(arrays[x]) - 1:
                continue
            heapq.heappush(heap, (arrays[x][y + 1], x, y + 1 ))
            
        return result