1329 将矩阵按对角线排序

题目描述#

矩阵对角线是一条从矩阵最上面行或者最左侧列中的某个元素开始的对角线，沿右下方向一直到矩阵末尾的元素。例如，矩阵 mat 有 6 行 3 列，从 mat[2][0] 开始的矩阵对角线将会经过 mat[2][0]、mat[3][1] 和 mat[4][2] 。

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

知识边界#

1. 矩阵对角线的数学性质#

在二维矩阵中，对于从左上到右下方向的同一条对角线上的任意元素 (i, j)，它们的行索引与列索引的差值是固定的，即 i - j 为一个常数。

i - j 的最小值为 0 - (n - 1) = 1 - n（右上角）
i - j 的最大值为 (m - 1) - 0 = m - 1（左下角）

利用这一特性，我们可以通过代数变换，将所有的对角线映射到一段连续的正整数区间内，从而只使用一个外层循环就可以遍历完所有的对角线。

2. Go 1.21 新特性#

泛型切片包 slices：标准库引入了 slices 包，使用 slices.Sort(slice) 可以对任意支持全序比较的泛型切片进行原地排序，语义更直观，性能也优于传统的 sort.Ints()。
内置的 min 和 max 函数：Go 1.21 起自带了内置的泛型 min 和 max 函数，无需再手动实现针对整型的比大小函数。

思路#

数学映射法#

不同于寻找对角线起点的模拟方式，我们可以利用对角线上 i - j 为常数的性质，直接对所有的对角线进行编号映射。

定义映射关系 $k$：已知 $i - j \in [1-n, m-1]$，为了让它变成一个方便遍历的正整数，我们可以加上偏移量 $m + n$。令 $k = i - j + m + n$。那么 $k$ 的范围就是 $[m + 1, 2m + n - 1]$。代码中从 $k = m$ 开始遍历到 $2m + n - 1$ 是完全可以覆盖所有有效对角线的。
反推列索引 $j$ 的边界：已知 $k = i - j + m + n$，可得 $i = k + j - m - n$。因为行索引 $i$ 必须满足 $0 \le i \le m - 1$，将其代入：
- $0 \le k + j - m - n \implies j \ge m + n - k$
- $k + j - m - n \le m - 1 \implies j \le 2m + n - 1 - k$
同时，列索引 $j$ 自身还需要满足矩阵边界 $0 \le j \le n - 1$。综合两部分限制条件，我们可以得出遍历 $j$ 时的合法区间：
- 最小值 minJ = max(0, m + n - k)
- 最大值 maxJ = min(n - 1, 2m + n - 1 - k)
执行排序：在明确了某一条对角线 $k$ 对应的 $j$ 的起止范围后，遍历 $j$，通过 $i = k + j - m - n$ 拿到对应元素并存入临时切片。用 slices.Sort() 对其进行排序后，再遍历一次放回原矩阵即可。

代码#

import "slices"

func diagonalSort(mat [][]int) [][]int {
	m := len(mat)
	n := len(mat[0])
	
	// 令 k = i - j + m + n
	// 对角线映射范围: k 的最大值在左下角，最小有效值在右上角
	for k := m; k < (2*m + n); k++ {
		maxJ := min(2*m+n-1-k, n-1)
		minJ := max(0, m+n-k)
		
		ans := []int{}
		
		// 收集对角线元素
		for j := minJ; j <= maxJ; j++ {
			ans = append(ans, mat[k+j-m-n][j])
		}
		
		// 使用 Go 1.21+ 的 slices.Sort 进行升序排序
		slices.Sort(ans)
		
		// 将排序好的元素写回原矩阵
		for j := minJ; j <= maxJ; j++ {
			mat[k+j-m-n][j] = ans[j-minJ]
		}
	}
	return mat
}