498 Diagonal Traverse

498. Diagonal Traverse

1. Question

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:

2. Implementation

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

        int m = matrix.length, n = matrix[0].length;
        int[] res = new int[m * n];
        int row = 0, col = 0;

        for (int i = 0; i < m * n; i++) {
            res[i] = matrix[row][col];
            // Moving up first
            if ((row + col) % 2 == 0) {
                if (col == n - 1) {
                    ++row;
                }
                else if (row == 0) {
                    ++col;
                }
                else {
                    --row;
                    ++col;
                }
            }
            else {
                if (row == m - 1) {
                    ++col;
                }
                else if (col == 0) {
                    ++row;
                }
                else {
                    ++row;
                    --col;
                }
            }
        }
        return res;
    }
}

3. Time & Space Complexity

时间复杂度O(mn), 空间复杂度O(mn)