中等

问题描述

给你一个大小为 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 <= 104
• 1 <= m * n <= 104
• -105 <= mat[i][j] <= 105

模拟

根据题意模拟，假设 $n$ 为矩阵行数，$m$ 为矩阵列数，$i$ 为当前遍历到的行数，$j$ 为当前遍历到的列数，遍历方向 $dir = 1$ 为向上，$dir = -1$ 为向下。可以发现，当方向是向上时， $i$ 递减，$j$ 递增，当方向是向下时， $i$ 递增，$j$ 递减。当到达矩阵边界时有以下几种情况：

• 上边界，$i = 0$，$j < m - 1$，且方向是向上时，需要往右移动一格，方向变为向下
• 右边界，$j = m - 1$，$i < n - 1$，且方向是向上时，需要往下移动一格，方向变为向下
• 下边界，$i = n - 1$，$j < m - 1$，且方向是向下时，需要往右移动一格，方向变为向上
• 左边界，$j = 0$，$i < n - 1$，且方向是向下时，需要往下移动一格，方向变为向上

根据以上规律进行遍历即可。

/**
 * @param {number[][]} mat
 * @return {number[]}
 */
var findDiagonalOrder = function (mat) {
  const n = mat.length
  const m = mat[0].length
  let i = 0
  let j = 0
  let dir = 1
  const ans = []
  while (i < n && j < m) {
    ans.push(mat[i][j])
    if (i === 0 && j < m - 1 && dir === 1) {
      j++
      dir = -1
    } else if (j === m - 1 && i < n - 1 && dir === 1) {
      i++
      dir = -1
    } else if (i === n - 1 && j < m - 1 && dir === -1) {
      j++
      dir = 1
    } else if (j === 0 && i < n - 1 && dir === -1) {
      i++
      dir = 1
    } else {
      j += dir
      i -= dir
    }
  }
  return ans
}

• 时间复杂度：$O(n \times m)$
• 空间复杂度：$O(1)$，返回值不算入额外的空间