[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]

解释:

说明:

1. 给定矩阵中的元素总数不会超过 100000 。

思路:

注意到，位置 (i, j) 在同一条对角线上 ⟺ i + j = 常数。没有注意到这一点的话，这题的循环写起来就很恶心了。那么循环顺序该怎么定呢？什么时候向上访问元素什么时候向下呢？i + j = 常数s：s为偶数向上，反之s为奇数则向下。于是可以得到大概的遍历代码骨架：

// 枚举每条对角线
for (s = 0; s <= m+n-2; s++) {      
    // 向上走
    if (s % 2 == 0) {                
        找合法的 i 范围;
        for (i 从大到小) 访问 mat[i][s-i];
    }
    // 向下走 
    else {                         
        找合法的 i 范围;  
        for (i 从小到大) 访问 mat[i][s-i];
    }
}Code language: C++ (cpp)

AC代码:

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        if (mat.empty() || mat[0].empty()) return {};
        
        int m = mat.size(), n = mat[0].size();
        vector<int> res;
        
        for (int s = 0; s < m + n - 1; s++) {
            int rStart = max(0, s - (n - 1));
            int rEnd = min(s, m - 1);
            
            if (s % 2 == 0) {
                for (int i = rEnd; i >= rStart; i--) {
                    res.push_back(mat[i][s - i]);
                }
            } else {
                for (int i = rStart; i <= rEnd; i++) {
                    res.push_back(mat[i][s - i]);
                }
            }
        }
        
        return res;
    }
};Code language: C++ (cpp)