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
from heapq import heappush, heappop
from typing import List


def sort_diagonally(mat: List[List[int]]) -> List[List[int]]:
    # If the input is a vector, return the vector
    if len(mat) == 1 or len(mat[0]) == 1:
        return mat

    # Rows + columns - 1
    # The -1 helps you to not repeat a column
    for i in range(len(mat)+len(mat[0])-1):
        # Process the rows
        if i+1 < len(mat):
            # Initialize heap, set row and column
            h = []
            row = len(mat)-(i+1)
            col = 0

            # Traverse diagonally, and add the values to the heap
            while row < len(mat):
                heappush(h, (mat[row][col]))
                row += 1
                col += 1

            # Sort the diagonal
            row = len(mat)-(i+1)
            col = 0
            while h:
                ele = heappop(h)
                mat[row][col] = ele
                row += 1
                col += 1
        else:
            # Process the columns
            # Initialize heap, row and column
            h = []
            row = 0
            col = i - (len(mat)-1)

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

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

    # Return the updated matrix
    return mat
