498. 对角线遍历

题目

给你一个大小为 m x n 的矩阵 mat ，请以对角线遍历的顺序，用一个数组返回这个矩阵中的所有元素。

难度：🌟🌟🌟🌟

示例 1：

输入：mat = [[1,2,3],[4,5,6],[7,8,9]]
输出：[1,2,4,7,5,3,6,8,9]

示例 2：

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

提示：

• m == mat.length
• n == mat[i].length
• 1 <= m, n <= 10⁴
• 1 <= m * n <= 10⁴
• -10⁵ <= mat[i][j] <= 10⁵

题解

模拟对角线遍历的过程，遍历方向由层数决定，而层数即为横纵坐标之和。

代码实现

func findDiagonalOrder(mat [][]int) []int {
	if mat == nil || len(mat) == 0 {
		return nil
	}

	m, n := len(mat), len(mat[0])
	row, col := 0, 0
	res := make([]int, m*n)
	for i := 0; i < len(res); i++ {
		res[i] = mat[row][col]
		// row+col 即为遍历的层数，偶数是向上遍历，奇数是向下遍历
		if ((row + col) % 2) == 0 { // 偶数
			if col == n-1 {
				// 往下移动一格，准备向下遍历
				row++
			} else if row == 0 {
				// 往右移动一格，准备向下遍历
				col++
			} else {
				// 往上移动
				row--
				col++
			}
		} else { // 奇数
			if row == m-1 {
				// 往右移动一格，准备向上遍历
				col++
			} else if col == 0 {
				// 往下移动一格，准备向上遍历
				row++
			} else {
				// 往下移动
				row++
				col--
			}
		}
	}
	return res
}