498. 对角线遍历

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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/diagonal-traverse
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

class Solution {
    public int[] findDiagonalOrder(int[][] mat) {
        if (mat == null || mat.length == 0 || mat[0].length == 0) {
            return new int[0];
        }
        int n = mat.length;
        int m = mat[0].length;
        int[] ans = new int[n * m];
        int index = 0;
        int row1 = 0, col1 = 0, row2 = 0, col2 = 0;
        boolean direct = true;

        while (row1 < n && col1 < m && row2 < n && col2 < m) {
            if (direct) {
                int x = row1;
                int y = col1;
                while (x >= row2 && y <= col2) {
                    ans[index++] = mat[x--][y++];
                }
            } else {
                int x = row2;
                int y = col2;
                while (x <= row1 && y >= col1) {
                    ans[index++] = mat[x++][y--];
                }
            }
            direct = !direct;

            if (row1 < n - 1) {
                row1++;
            } else {
                col1++;
            }

            if (col2 < m - 1) {
                col2++;
            } else {
                row2++;
            }

        }
        return ans;
    }
}