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:

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

public class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix.length == 0 || matrix[0].length == 0)
            return new int[0];
        
        int m = matrix.length, n = matrix[0].length, N = m * n;
        int[] result = new int[N];
        int i = 0, j = 0, index = 0, time = 1;
        while (index < N && i < m && j < n) {
            if (time == 1) {
                while(i >= 0 && j < n) {
                    result[index++] = matrix[i][j];
                    
                    i -= 1;
                    j += 1;
                    
                    if (j == n) 
                        i += 2;
                }
                i = Math.max(0, i);
                j = Math.min(n - 1, j);
                
                time = 2;
            } else {
                while (i < m && j >= 0) {
                    result[index++] = matrix[i][j];
                    i += 1;
                    j -= 1;
                    
                    if (i == m)
                        j += 2;
                }
                i = Math.min(m - 1, i);
                j = Math.max(0, j);
                
                time = 1;
            }
        }
        return result;
    }
}