[LeetCode]498. 对角线遍历

题目

498. 对角线遍历
给你一个大小为 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

public int[] findDiagonalOrder(int[][] matrix) {
    if (matrix == null || matrix.length == 0 || matrix[0].length == 0) return new int[]{};
    int R = matrix.length, C = matrix[0].length;
    int[] ans = new int[R * C];
    int r = 0, c = 0, idx = 0;
    for (int i = 0; i < R + C - 1; i++) {
        if (i % 2 == 0) {
            while (r >= 0 && c < C) {
                ans[idx++] = matrix[r][c];
                r--;
                c++;
            }
            if (c < C) {
                r++;
            } else {
                r += 2;
                c--;
            }

        } else {
            while (r < R && c >= 0) {
                ans[idx++] = matrix[r][c];
                r++;
                c--;
            }
            if (r < R) {
                c++;
            } else {
                c += 2;
                r--;
            }
        }
    }
    return ans;
}