Diagonal Traverse II LeetCode solution in Python

1
class Solution:
2
    def findDiagonalOrder(self, nums: List[List[int]]) -> List[int]:
3
        # sort the index of element
4
        heap = list()
5
        n = len(nums)
6
        for i in range(n):
7
            m = len(nums[i])
8
            for j in range(m):
9
                heapq.heappush(heap, [i + j, j, i])
10

11
        # append the element one by one
12
        ans = []
13
        while heap:
14
            temp = heapq.heappop(heap)
15
            ans.append(nums[temp[2]][temp[1]])
16
        return ans