498. 对角线遍历（模拟）

给你一个大小为 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

// 使用v与h控制上下与左右的前进方向，分别处理四个对角的特殊情况即可
class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        int v = -1;
        int h = 1;
        int i = 0, j = 0;
        vector<int> res;
        while (1) {
            res.push_back(mat[i][j]);
            if (i == mat.size() - 1 && j == mat[0].size() - 1)
                break;
            if (i + v < 0) {
                if (j + h >= mat[0].size())
                    i += 1;
                else
                    j += 1;
                v = 1;
                h = -1;
            }
            else if (j + h < 0) {
                if (i + v >= mat.size())
                    j += 1;
                else
                    i += 1;
                v = -1;
                h = 1;
            }
            else if (j + h >= mat[0].size()) {
                i += 1;
                v = 1;
                h = -1;
            }
            else if (i + v >= mat.size()) {
                j += 1;
                v = -1;
                h = 1;
            }
            else {
                i += v;
                j += h;
            }
        }
        return res;
    }
};