LeetCode498. 对角线遍历(C++)

给你一个大小为 m x n 的矩阵 mat ，请以对角线遍历的顺序，用一个数组返回这个矩阵中的所有元素。

 

示例 1：
输入：mat = [[1,2,3],[4,5,6],[7,8,9]]
输出：[1,2,4,7,5,3,6,8,9]
示例 2：

输入：mat = [[1,2],[3,4]]
输出：[1,2,3,4]
 

提示：

m == mat.length
n == mat[i].length
1 <= m, n <= 104
1 <= m * n <= 104
-105 <= mat[i][j] <= 105

思路

一共共有 m + n - 1条对角线，我们从0开始编号

可以分为两类

1.从上到下的

此时序号为奇数(i % 2 == 1)

若i < n 则起始坐标为 [0, i]

若 i >= n 则起始坐标为 [i - n + 1, n - 1]

2.从下到上的

此时序号为偶数(i % 2 == 0)

若i < m 则起始坐标为 [i, 0]

若 i >= m 则起始坐标为 [m - 1, i - m + 1]

代码

class Solution {
public:
    vector<int> findDiagonalOrder(vectorint>>& mat) {
        int m = mat.size();
        int n = mat[0].size();
        vector<int> res;
 
        for(int i = 0; i < m + n - 1; i ++) //一共有n + m - 1条对角线
        {
            if(i % 2)//奇数，说明是从上往下
            {   
                int x = i < n ? 0 : i - n + 1;
                int y = i < n ? i : n - 1;
                while(x < m && y >= 0)
                {
                res.push_back(mat[x][y]);
                x ++, y --;
                }
            }
            else //偶数，说明是从下往上
            {
                int x = i < m ? i : m - 1;
                int y = i < m ? 0 : i - m + 1;
                while(x >= 0 && y < n)
                {
                    res.push_back(mat[x][y]);
                    x --, y ++;
                }
            }
        }
        return res;
    }
};