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:

ex1

Note:

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

Solution:

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

        int[] list = new int[rows * cols];
        int row = 0;
        int col = 0;
        int direction = 1;
        
        for (int i = 0; i < list.length; i++) {
            list[i] = A[row][col];
            if (direction == 1) {
                if (col == cols - 1) {
                    direction = -1;
                    row++;                    
                } else if (row == 0) {
                    col++;
                    direction = -1;
                } else {
                    row--;
                    col++; 
                }
            } else {
                if (row == rows - 1) {
                    col++;
                    direction = 1;
                } else if (col == 0) {
                    row++;
                    direction = 1;
                } else {
                    row++;
                    col--;
                }
            }
        }
        return list;
    }
}