Merge K Sorted Array

Merge K sorted array into one big sorted array in ascending order.

Assumptions

The input arrayOfArrays is not null, none of the arrays is null either.

import heapq
class Solution(object):
  def merge(self, arrayOfArrays):
    heap = []
    for i in xrange(len(arrayOfArrays)):
      if len(arrayOfArrays[i]):
        heap.append((arrayOfArrays[i][0],i,0))
    heapq.heapify(heap)
    res = []
    while heap:
      val,index,element = heapq.heappop(heap)
      res.append(val)
      if element + 1 < len(arrayOfArrays[index]):
        heapq.heappush(heap,(arrayOfArrays[index][element + 1],index,element + 1))
    return res