heapq in Python to print all elements in sorted order from row and column wise sorted matrix

Given an n x n matrix, where every row and column is sorted in non-decreasing order. Print all elements of matrix in sorted order. Examples:

Input : mat= [[10, 20, 30, 40],
              [15, 25, 35, 45],
              [27, 29, 37, 48],
              [32, 33, 39, 50]]

Output : Elements of matrix in sorted order
         [10, 15, 20, 25, 27, 29, 30, 32, 
          33, 35, 37, 39, 40, 45, 48, 50]

This problem has existing solution please refer link. We will solve this problem in python with the same approach of merging two sorted arrays using heapq module. 

Implementation:

Python3

# Function to print all elements in sorted order
# from row and column wise sorted matrix
from heapq import merge

def sortedMatrix(mat):

    # initialize result variable with first row of matrix
    result=mat[0]

    # now traverse through complete matrix
    # after first row and merge each row with
    # result one by one
    # after last operation result will contain
    # list of sorted elements of matrix
    for row in mat[1:]:
        result=list(merge(result,row))

    return result

if __name__ == "__main__":
    mat = [[10, 20, 30, 40],
        [15, 25, 35, 45],
        [27, 29, 37, 48],
        [32, 33, 39, 50]]
    print ('Elements of matrix in sorted order')
    print (sortedMatrix(mat))

Output

Elements of matrix in sorted order
[10, 15, 20, 25, 27, 29, 30, 32, 33, 35, 37, 39, 40, 45, 48, 50]