Diagonal Traverse

Problem

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.

Solution

public class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix == null || matrix.length == 0) return new int[0];
        int row = matrix.length, col = matrix[0].length;
        int[] res = new int[row * col];
        int m = 0, n = 0, direction = 1;
        for (int i = 0; i < res.length; i++) {
            res[i] = matrix[m][n]; 
            m -= direction;
            n += direction;
            //Handle overloaded corner cases 
            //Must put first two if statements at the beginning
            if (m >= row) {
                m = row - 1;
                n += 2;
                direction *= -1;
            }
            if (n >= col) {
                n = col - 1;
                m += 2;
                direction *= -1;
            }
            if (m < 0) {
                m = 0;
                direction *= -1;
            }
            if (n < 0) {
                n = 0;
                direction *= -1;
            }
        }
        return res;
    }
}

Analysis

The problem looks tricky at first, however, as long as we find the pattern, it becomes much easier
The pattern here is that every time when traversal exceeds matrix's border, it changes its direction
And the direction decides how row and col index will be changed
Hence we have direction = 1 and update it by direction *= -1 when we exceeds the border
Notice that we have to check m >= row and n >= col first
Otherwise we will mistakenly change direction twice because when m >= row, n will be < 0 as well