[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 表示是往右上还是往左下。起点是mat[0][0]。唯一需要注意的是往右上走的时候，当我们碰到矩阵的边界的时候，应该先判断 col 再判断 row。因为如果我们在第一行，row=0，col++，那么 row 还是 0，所以我们应该先判断 col 是否到达边界。往左下走的时候，我们应该先判断 row 再判断 col，因为如果我们在第一列，col=0，row++，那么 col 还是 0，所以我们应该先判断 row 是否到达边界。

复杂度

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

代码

Java实现

class Solution {
    public int[] findDiagonalOrder(int[][] mat) {
        int row = 0;
        int col = 0;
        int m = mat.length;
        int n = mat[0].length;
        int dir = 1;
        int[] res = new int[m * n];
        for (int i = 0; i < res.length; i++) {
            res[i] = mat[row][col];
            // up
            if (dir == 1) {
                if (col == n - 1) {
                    row++;
                    dir = -1;
                } else if (row == 0) {
                    col++;
                    dir = -1;
                } else {
                    row--;
                    col++;
                }
            }
            // down
            else {
                if (row == m - 1) {
                    col++;
                    dir = 1;
                } else if (col == 0) {
                    row++;
                    dir = 1;
                } else {
                    row++;
                    col--;
                }
            }
        }
        return res;
    }
}