LeetCode498. 对角线遍历

题目描述

本题目来自LeetCode上的『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]

提示

• m == mat.length
• n == mat[i].length
• 1 <= m, n <= 10^4
• 1 <= m * n <= 10^4
• -10^5 <= mat[i][j] <= 10^5

题解

设行数有 $m$ 行，列数有 $n$ 列，则一共有 $m+n-1$ 个对角线，编号为 $[0,m+n-2]$，当编号为偶数时，从左下往右上遍历；当编号为奇数时，从右上往左下遍历。

代码

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        int m = mat.size(), n = mat[0].size();
        vector<int> ans;
        for (int i = 0; i < m + n - 1; ++i) {
            if (i & 1) {
                int x = i < n ? 0 : i - n + 1;
                int y = i - x;
                while (x < m && y >= 0) {
                    ans.emplace_back(mat[x++][y--]);
                }
            } else {
                int y = i < m ? 0 : i - m + 1;
                int x = i - y;
                while (x >= 0 && y < n) {
                    ans.emplace_back(mat[x--][y++]);
                }
            }
        }
        return ans;
    }
};

复杂度分析

• 时间复杂度：$O(mn)$
• 空间复杂度：$O(1)$