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

闪念标签:LC

题目链接:https://leetcode.cn/problems/sort-the-matrix-diagonally/