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<int> findDiagonalOrder(vector<vector<int>>& matrix) {
if(matrix.size() == 0 || matrix[0].size() == 0) return vector<int>();
int direction = 0;
int M = static_cast<int>(matrix.size()), N = static_cast<int>(matrix[0].size());
vector<int> 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;
}
};