498. 对角线遍历

难度：中等

给定一个含有 M x N 个元素的矩阵（M 行，N 列），请以对角线遍历的顺序返回这个矩阵中的所有元素，对角线遍历如下图所示。

 

示例:

输入:
[
 [ 1, 2, 3 ],
 [ 4, 5, 6 ],
 [ 7, 8, 9 ]
]

输出:  [1,2,4,7,5,3,6,8,9]

解释:

说明:

1. 给定矩阵中的元素总数不会超过 100000 。

分析：

模拟，↗或者↙两个方向，出界了调整位置

代码：

public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix == null || matrix.length == 0) {
            return new int[0];
        }
        int N = matrix.length;
        int M = matrix[0].length;

        int row = 0, column = 0;

        int direction = 1;

        int[] result = new int[N*M];
        int index = 0;

        while (row < N && column < M) {

            result[index++] = matrix[row][column];
            //direction==1 则往↗方向
            int new_row = row + (direction == 1 ? -1 : 1);
            int new_column = column + (direction == 1 ? 1 : -1);

            //如果出界了
            if (new_row < 0 || new_row == N || new_column < 0 || new_column == M) {
                if (direction == 1) {
                    row += (column == M - 1 ? 1 : 0) ;
                    column += (column < M - 1 ? 1 : 0);
                } else {
                    column += (row == N - 1 ? 1 : 0);
                    row += (row < N - 1 ? 1 : 0);
                }

                direction = 1 - direction;

            } else {
                row = new_row;
                column = new_column;
            }
        }
        return result;
    }

结果：