题目
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;
}