Sort the diagonals in ascending order from bottom left to top right of a given a squared matrix of positive integers implemented in python

sort_diagonal.py

def solution(matrix):
    M, N, diagonals = len(matrix), len(matrix[0]), []

    # 1. Get all the diagonals
    for row in range(0, M):
        col, temp, i = 0, [], row
        j = col
        
        while i >= 0 and j < N:
            temp.append(matrix[i][j])
            i -= 1
            j += 1

        # 2. Sort the diagonals
        temp.sort()
        diagonals.append(temp)

    for col in range(1, N):
        row, temp, i, j = M - 1, [], row, col
        
        while i >= 0 and j < N:
            temp.append(matrix[i][j])
            i -= 1
            j += 1
        
        # 2. Sort the diagonals
        temp.sort()
        diagonals.append(temp)

    # 3. Rebuild the final matrix
    finalMatrix = []
    for row in range(0, M):
        i = row + 1
        for j in range(1, N):
            current_row = diagonals[row]
            next_row = diagonals[i]
            el = next_row.pop()
            current_row.append(el)
            i += 1
        finalMatrix.append(current_row)
    
    return print(finalMatrix)


# test case 03
mat3 = [
  [1, 3, 9, 4],
  [9, 5, 7, 7],
  [3, 6, 9, 7],
  [1, 2, 2, 2],
];

solution(mat3)

# Expected output
# [
#   [1, 9, 9, 7],
#   [3, 5, 6, 9],
#   [3, 4, 7, 7],
#   [1, 2, 2, 2],
# ];