[LeetCode] 498. Diagonal Traverse

Given an m x n matrix mat, return an array of all the elements of the array in a diagonal order.

Example 1:

Input: mat = [[1,2,3],[4,5,6],[7,8,9]]
Output: [1,2,4,7,5,3,6,8,9]

Example 2:
Input: mat = [[1,2],[3,4]]
Output: [1,2,3,4]

Constraints:
m == mat.length
n == mat[i].length
1 <= m, n <= 104
1 <= m * n <= 104
-105 <= mat[i][j] <= 105

对角线遍历。

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

思路

这个题不涉及任何算法，难点是如何判断到底是往上扫描还是往下扫描以及需要换方向的时候的一些相关操作。具体思路是

• 设一个变量 dir 表示到底是往右上还是左下

• 当 dir == 1（向右上 ↗️）时，下一步是 row - 1, col + 1。到达右上角 (row == 0 && col == n-1) 时：

  • 既满足 row == 0，也满足 col == n - 1。
  • 正确做法是优先判断右边界 col == n - 1，然后下移 row++ 并切换方向；
  • 如果先判断上边界 row == 0，你会去执行 col++，直接越界。

• 同理，当 dir == -1（向左下 ↙️）时，下一步是 row + 1, col - 1。到达左下角 (row == m - 1 && col == 0) 时：

  • 既满足 row == m - 1，也满足 col == 0。

• 正确做法是优先判断下边界 row == m - 1，然后右移 col++ 并切换方向；

• 如果先判断左边界 col == 0，你会去执行 col–，直接越界。

复杂度

时间O(mn)
空间O(n) - 存储output

代码

Java实现

class Solution {
	public int[] findDiagonalOrder(int[][] matrix) {
		// corner case
		if (matrix == null || matrix.length == 0) {
			return new int[0];
		}

		// normal case
		int row = 0;
		int col = 0;
		int m = matrix.length;
		int n = matrix[0].length;
		int[] res = new int[m * n];
		for (int i = 0; i < res.length; i++) {
			res[i] = matrix[row][col];
			// moving up
			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;
	}
}