JZX轻语:简
LeetCode 1329 - 将矩阵按对角线排序
发表于2024年04月29日
不难得知,在每个对角线的元素中,列号和行号的差值是恒定的,可以根据行号确定某个对角线元素的列号。按行/列其中一个维度斜着就地快速排序即可,需要注意边界条件。
class Solution:
def diagonalSort(self, mat: List[List[int]]) -> List[List[int]]:
def partition(lo: int, hi: int, diff: int):
nonlocal mat
pivot = mat[hi][hi + diff]
i = lo
for j in range(lo, hi):
if mat[j][j + diff] < pivot:
mat[i][i + diff], mat[j][j + diff] = mat[j][j + diff], mat[i][i + diff]
i += 1
mat[i][i + diff], mat[hi][hi + diff] = mat[hi][hi + diff], mat[i][i + diff]
return i
def qsort(lo: int, hi: int, diff: int):
""" 针对mat[lo][lo+diff], mat[lo+1][lo+1+diff], ..., mat[hi][hi+diff]进行排序 """
nonlocal mat
if lo >= hi:
return
pivot_pos = partition(lo, hi, diff)
qsort(lo, pivot_pos - 1, diff)
qsort(pivot_pos + 1, hi, diff)
m = len(mat)
n = len(mat[0])
for i in range(0, n):
qsort(0, min(m - 1, n - 1 - i), i)
for i in range(1, m):
qsort(i, min(m - 1, n - 1 + i), -i)
return mat