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:

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

vector<int> findDiagonalOrder(vector<vector<int>>& matrix) { // time: O(m * n); space: O(m * n)
    if (matrix.empty() || matrix[0].empty()) return {};
    int m = matrix.size(), n = matrix[0].size(), r = 0, c = 0;
    vector<int> res(m * n);
    for (int i = 0; i < m * n; ++i) {
        res[i] = matrix[r][c];
        if ((r + c) % 2 == 0) {
            if (c == n - 1) ++r;
            else if (r == 0) ++c;
            else { // move to upper right
                --r, ++c;
            }
        } else {
            if (r == m - 1) ++c;
            else if (c == 0) ++r;
            else { // move to bottom left
                ++r, --c;
            }
        }
    }
    return res;
}

vector<int> findDiagonalOrder(vector<vector<int>>& matrix) { // time: O(m * n); space: O(m * n)
    if (matrix.empty() || matrix[0].empty()) return {};
    int m = matrix.size(), n = matrix[0].size(), k = 0;
    vector<int> res(m * n);
    for (int i = 0; i < m + n - 1; ++i) {
        int r_low = max(0, i - n + 1), r_high = min(i, m - 1);
        if (i % 2 == 0) {
            for (int r = r_high; r >= r_low; --r) {
                res[k++] = matrix[r][i - r];
            }
        } else {
            for (int r = r_low; r <= r_high; ++r) {
                res[k++] = matrix[r][i - r];
            }
        }
    }
    return res;
}