498. Diagonal Traverse

Given a matrix of M x N elements (M rows, N columns), return all elements of the matrix in diagonal order as shown in the below image.

Example:

Input:
[
[ 1, 2, 3 ],
[ 4, 5, 6 ],
[ 7, 8, 9 ]
]
Output:  [1,2,4,7,5,3,6,8,9]

Explanation:



Note:

The total number of elements of the given matrix will not exceed 10,000.

class Solution {
public int[] findDiagonalOrder(int[][] matrix) {
if(matrix == null || matrix.length == 0 || matrix[0] == null || matrix[0].length == 0)
return new int[0];

int m = matrix.length, n = matrix[0].length;
int[] result = new int[m * n];
int i = 0;
for(int count = 0; count <= m + n; count++) {
if(count % 2 == 0) {
for(int j = Math.max(0, count - m + 1); j <= Math.min(count, n - 1); j++) {
result[i++] = matrix[count - j][j];
}
} else {
for(int j = Math.max(0, count - n + 1); j <= Math.min(count, m - 1); j++) {
result[i++] = matrix[j][count - j];
}
}
}

return result;
}
}