leetcode 链接：498. 对角线遍历

题目

给定一个含有 M x N 个元素的矩阵（M 行，N 列），请以对角线遍历的顺序返回这个矩阵中的所有元素，对角线遍历如下图所示。

示例：

输入:
[
 [ 1, 2, 3 ],
 [ 4, 5, 6 ],
 [ 7, 8, 9 ]
]
输出:  [1,2,4,7,5,3,6,8,9]

解答 & 代码

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        vector<int> result;
        if(mat.size() == 0 || mat[0].size() == 0)
            return result;
        int rows = mat.size();
        int cols = mat[0].size();

        int row = 0;        // 当前走到的行号
        int col = 0;        // 当前走到的列号
        int rowStep = -1;    // 每一步行号的更新
        int colStep = 1;    // 每一步列号的更新
        // 遍历对角线
        while(row != rows - 1 || col != cols - 1)
        {
            // 如果当前行号、列好都没越界，说明还在正常地遍历当前对角线
            // 则将当前位置的元素压入结果数组，更新行号、列号（即继续遍历对角线）
            if(row >= 0 && row < rows && col >= 0 && col < cols)
            {
                result.push_back(mat[row][col]);
                row += rowStep;
                col += colStep;
            }
            else    // 越界
            {
                // 注意先处理右边界再处理上边界，先处理下边界再处理左边界
                // 因为会有同时满足超出右边界和上边界的情况，应该按右边界越界处理
                // 也会有同时满足超出下边界和左边界的情况，应该按下边界越界处理
                // 如果超过右边界
                if(col >= cols)
                {
                    row += 2;
                    col -= 1;
                }
                // 如果超过上边界
                else if(row < 0)
                    ++row;
                // 如果超过下边界
                else if(row >= rows)
                {
                    col += 2;
                    row -= 1;
                }
                // 如果超过左边界
                else if(col < 0)
                    ++col;

                // 接下来要遍历另外的对角线，因此改变更新步伐
                rowStep = -rowStep;
                colStep = -colStep;
            }
        }
        result.push_back(mat[row][col]);
        return result;
    }
};

执行结果：

执行结果：通过

执行用时：28 ms, 在所有 C++ 提交中击败了 80.40% 的用户
内存消耗：17.9 MB, 在所有 C++ 提交中击败了 62.09% 的用户