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]\

Note:

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

Tips:

设置一个方向dir,如果碰到四边的壁就转向。注意转向时要先判断是否超出m和n，在判断是否小于0，不然会出错。

O(n)

Code：

public class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix == null || matrix.length == 0) return new int[0];
        int[] res = new int[matrix.length * matrix[0].length];
        int m = matrix.length, n = matrix[0].length;
        int col = 0, row = 0, d = 0;
        int[][] dir = {{-1, 1},{1, -1}};

        for (int i = 0; i < m * n; i++) {
            res[i] = matrix[row][col];
            row = row + dir[d][0];
            col = col + dir[d][1];
            if (row >= m) {
                row = m - 1;
                col += 2;
                d = 1 - d;
            }
            if (col >= n) {
                col = n - 1;
                row += 2;
                d = 1 - d;
            }
            if (row < 0) {
                row = 0;
                d = 1 - d;
            } 
            if (col < 0) {
                col = 0;
                d = 1 - d;
            } 
        }
        return res;
    }
}