将矩阵按对角线排序

1329. 将矩阵按对角线排序

给你一个 m * n 的整数矩阵 mat ，请你将同一条对角线上的元素（从左上到右下）按升序排序后，返回排好序的矩阵。

示例 1：

输入：mat = [[3,3,1,1],[2,2,1,2],[1,1,1,2]]

输出：[[1,1,1,1],[1,2,2,2],[1,2,3,3]]

def func(mat):
    m, n = len(mat), len(mat[0])
    start_point = [(i, 0) for i in range(0, m)] + [(0, j) for j in range(1, n)]
    for tup in start_point:
        i, j = tup
        res = [mat[i + k][j + k] for k in range(0, min(m, n)) if i + k < m and j + k < n]
        res.sort()
        for k in range(len(res)):
            mat[i + k][j + k] = res[k]
    return mat


print(func([[3,3,1,1],[2,2,1,2],[1,1,1,2]]))