498. Diagonal Traverse

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.

--------------------------
分（上，右）和（下，左）2种情况，长方形的4个顶点可能需要特殊处理

class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix.length == 0) return new int[0];
        
        int m = matrix.length, n = matrix[0].length;
        int[] rt = new int[m * n];
        int row = 0, col = 0, i = 0;
        int up = 1;
        
        while (row != m - 1 || col != n - 1) {
            rt[i] = matrix[row][col];
            
            if ((row == 0 || col == n - 1) && up == 1) {
                if (col != n - 1 && row == 0) {
                    col++;
                }else {
                    row++;
                }
                
                up = -1;
            }else if ((row == m - 1 || col == 0) && up == -1) {
                if (row == m - 1) {
                    col++;
                }else {
                    row++;
                }
                
                up = 1;
            } else {
                row -= up;
                col += up;
            }
            i++;
        }
        
        rt[m * n - 1] = matrix[m - 1][n - 1];
        return rt;
    }
}