clueless coding // TODO: be smarter

LeetCode 498. Diagonal Traverse: C++ Solution



Problem Statement


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.

Function Signature (C++):
  vector<int> findDiagonalOrder(vector<vector<int>>& matrix)

Inputs:
[
 [ 1, 2, 3 ],
 [ 4, 5, 6 ],
 [ 7, 8, 9 ]
]

Outputs:
  [1,2,4,7,5,3,6,8,9]



TL;DR Code Solution


vector<int> findDiagonalOrder(vector<vector<int>>& matrix) {

    vector<int> result;
    if (matrix.size() == 0 || matrix[0].size() == 0) {
        return result;
    }
    int m = matrix.size();
    int n = matrix[0].size();
    vector<vector<int>> dirs = { {-1, 1}, {1, -1} };

    int row = 0; int col = 0; int d = 0;

    for (int i = 0; i < m * n; i++) {
        result.push_back(matrix[row][col]);
        row += dirs[d][0];
        col += dirs[d][1];

        if (row >= m) { row = m - 1; col += 2; d = 1 - d;}
        if (col >= n) { col = n - 1; row += 2; d = 1 - d;}
        if (row < 0)  { row = 0; d = 1 - d;}
        if (col < 0)  { col = 0; d = 1 - d;}

    }

    return result;

}