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:

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

Subscribe to see which companies asked this question.

My solution:

public class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix == null || matrix.length == 0) return new int[0];
        int m = matrix.length;
		int n = matrix[0].length;
		int[] re = new int[m * n];
		boolean flag = true;
		int mnow = 0;
		int nnow = 0;
		for (int i = 0; i < m * n; ++i) {
			re[i] = matrix[mnow][nnow];
			if (flag) {
				if (mnow == 0 || nnow == n - 1) {
					if (nnow == n - 1)
						mnow++;
					else
						nnow++;
					flag = false;
				} else {
					mnow--;
					nnow++;
				}
				
			} else {
				if (mnow == m - 1 || nnow == 0) {
					if (mnow == m - 1)
						nnow++;
					else
						mnow++;
					flag = true;
				} else {
					mnow++;
					nnow--;
				}
			}
		}
		return re;
    }
}