498. Diagonal Traverse (Medium)

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:

1. The total number of elements of the given matrix will not exceed 10,000. Show Company Tags

Solution 1: 96ms

I don't think this is a hard problem. It is easy to figure out the walk pattern. Anyway... Walk patterns:

• if out of bottom border (row >= m) then row = m - 1; col += 2; change walk direction.
• if out of right border (col >= n) then col = n - 1; row += 2; change walk direction.
• if out of top border (row < 0) then row = 0; change walk direction.
• if out of left border (col < 0) then col = 0; change walk direction.
• Otherwise, just go along with the current direction.

Time complexity: $$O(m * n)$$, m = number of rows, n = number of columns. Space complexity: $$O(1)$$

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& matrix) {
        if (matrix.empty() || matrix[0].empty()) return {};
        int m = matrix.size(), n = matrix[0].size();
        int r = 0, c = 0, d = 0; // first go up
        vector<int> res(m*n);
        vector<vector<int>> dirs = {{-1,1}, {1,-1}};
        for (int i = 0; i < m*n; ++i) {
            // cout << r << " " << c << " " << matrix[r][c] << endl;
            res[i] = matrix[r][c];
            r += dirs[d][0], c += dirs[d][1];
            if (r >= m) { r = m-1, c += 2, d = 1-d; }
            if (c >= n) { c = n-1, r += 2, d = 1-d; }
            if (r < 0) { r = 0, d = 1-d; }
            if (c < 0) { c = 0, d = 1-d; }
        }
        return res;
    }
};

Solution 2: 86ms

Put all diagonal sequences from top-right to bottom-left to an array and then combine all sequence together by reversing odd sequences.

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& matrix) {
        if (matrix.empty() || matrix[0].empty()) return {};
        int m = matrix.size(), n = matrix[0].size();
        bool reverse = true;
        vector<int> res;
        for (int i = 0; i < m+n-1; ++i) {
            vector<int> tmp;
            int r = max(0, i-n+1), c = min(i, n-1);
            while (r < m && c >= 0) {
                tmp.push_back(matrix[r++][c--]);
            }
            if (reverse) {
                res.insert(res.end(), tmp.rbegin(), tmp.rend());
            } else {
                res.insert(res.end(), tmp.begin(), tmp.end());
            }
            reverse = !reverse;
        }
        return res;
    }
};