given a m n matrix mat of integers sort it diagonally in ascending order from the topleft to the bottomright then return the sorted array mat 3 3 1 1 2 2 1 2 1 1 1 2 should return 1 1 1 1 1 2 2 2 1 2 3 3 if the input is a vector return the vector rows columns 1 the 1 helps you to not repeat a column process the rows initialize heap set row and column traverse diagonally and add the values to the heap sort the diagonal process the columns initialize heap row and column traverse diagonally sort the diagonal return the updated matrix if the input is a vector return the vector rows columns 1 the 1 helps you to not repeat a column process the rows initialize heap set row and column traverse diagonally and add the values to the heap sort the diagonal process the columns initialize heap row and column traverse diagonally sort the diagonal return the updated matrix
from heapq import heappush, heappop
from typing import List


def sort_diagonally(mat: List[List[int]]) -> List[List[int]]:
    
    if len(mat) == 1 or len(mat[0]) == 1:
        return mat

    
    
    for i in range(len(mat)+len(mat[0])-1):
        
        if i+1 < len(mat):
            
            h = []
            row = len(mat)-(i+1)
            col = 0

            
            while row < len(mat):
                heappush(h, (mat[row][col]))
                row += 1
                col += 1

            
            row = len(mat)-(i+1)
            col = 0
            while h:
                ele = heappop(h)
                mat[row][col] = ele
                row += 1
                col += 1
        else:
            
            
            h = []
            row = 0
            col = i - (len(mat)-1)

            
            while col < len(mat[0]) and row < len(mat):
                heappush(h, (mat[row][col]))
                row += 1
                col += 1

            
            row = 0
            col = i - (len(mat)-1)
            while h:
                ele = heappop(h)
                mat[row][col] = ele
                row += 1
                col += 1

    
    return mat