LeetCode498. 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:

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

class Solution {
public:
    vector findDiagonalOrder(vector>& matrix) {
        if(matrix.size() == 0 || matrix[0].size() == 0) return vector();
        int direction = 0;
        int M = static_cast(matrix.size()), N = static_cast(matrix[0].size());
        vector ans(M*N);
        int row_idx = 0, col_idx = 0, count = 0;
        while(count < M*N) {
            ans[count] = matrix[row_idx][col_idx];
            bool on_top_board = row_idx == 0;
            bool on_right_board = col_idx == N-1;
            bool on_left_board = col_idx == 0;
            bool on_bottom_board = row_idx == M-1;

            if(direction == 0) {
                if(on_right_board) {
                    ++row_idx;
                    direction = 1;
                }
                else if(on_top_board) {
                    ++col_idx;
                    direction = 1;
                }
                else {
                    --row_idx;
                    ++col_idx;
                }
            }
            else if(direction == 1) {
                if(on_bottom_board) {
                    ++col_idx;
                    direction = 0;
                }
                else if(on_left_board) {
                    ++row_idx;
                    direction = 0;
                }
                else {
                    ++row_idx;
                    --col_idx;
                }
            }
            count ++;
        }
        return ans;
    }
};