LC 498. 对角线遍历

题目描述

这是 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]

示例 2：

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

输出：[1,2,3,4]

提示：

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

---

模拟

根据题意进行模拟即可。

为了方便，令 mat 为 g，记 g 的行和宽分别为 $n$ 和 $m$。当前所在位置为 $(x, y)$，遍历方向使用 $dir$ 代指（当 $dir = 1$ 代表往右上方进行遍历，当 $dir = -1$ 代表往左下方进行遍历），使用 $idx$ 记录当前处理到的答案下标。

每次除了将当前格子放入答案（ans[idx++]=g[x][y]）以外，还需要结合 $dir$ 找到当前位置的右上方格子 $(x - 1, y + 1)$ 或是左下方格子 $(x + 1, y - 1)$，若下一目标位置「越界」并且还没搜索完整个矩阵，我们需要根据优先级来找「下一个发起点」的位置，并且翻转遍历方向。

具体的找「下一个发起点」的优先级为：

• 若当前遍历方向为往右上角，即 $dir = 1$，优先找 $(x, y + 1)$ 作为下一发起点，若越界，则找 $(x + 1, y)$ 作为下一发起点；
• 若当前遍历方向为往左下角，即 $dir = -1$，优先找 $(x + 1, y)$ 作为下一发起点，若越界，则找 $(x, y + 1)$ 作为下一发起点。

Java 代码：

class Solution {
    public int[] findDiagonalOrder(int[][] g) {
        int n = g.length, m = g[0].length, cnt = n * m;
        int[] ans = new int[cnt];
        int x = 0, y = 0, dir = 1, idx = 0;
        while (idx != cnt) {
            ans[idx++] = g[x][y];
            int nx = x, ny = y;
            if (dir == 1) {
                nx = x - 1; ny = y + 1;
            } else {
                nx = x + 1; ny = y - 1;
            }
            if (idx < cnt && (nx < 0 || nx >= n || ny < 0 || ny >= m)) {
                if (dir == 1) {
                    nx = y + 1 < m ? x : x + 1;
                    ny = y + 1 < m ? y + 1 : y;
                } else {
                    nx = x + 1 < n ? x + 1 : x;
                    ny = x + 1 < n ? y : y + 1;
                }
                dir *= -1;
            }
            x = nx; y = ny;
        }
        return ans;
    }
}

C++ 代码：

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& g) {
        int n = g.size(), m = g[0].size(), cnt = n * m;
        vector<int> ans(cnt);
        int x = 0, y = 0, dir = 1, idx = 0;
        while (idx != cnt) {
            ans[idx++] = g[x][y];
            int nx = x, ny = y;
            if (dir == 1) {
                nx = x - 1; ny = y + 1;
            } else {
                nx = x + 1; ny = y - 1;
            }
            if (idx < cnt && (nx < 0 || nx >= n || ny < 0 || ny >= m)) {
                if (dir == 1) {
                    nx = y + 1 < m ? x : x + 1;
                    ny = y + 1 < m ? y + 1 : y;
                } else {
                    nx = x + 1 < n ? x + 1 : x;
                    ny = x + 1 < n ? y : y + 1;
                }
                dir *= -1;
            }
            x = nx; y = ny;
        }
        return ans;
    }
};

Python 代码：

class Solution:
    def findDiagonalOrder(self, g: List[List[int]]) -> List[int]:
        n, m = len(g), len(g[0])
        cnt = n * m
        ans = [0] * cnt
        x, y, di, idx = 0, 0, 1, 0
        while idx != cnt:
            ans[idx] = g[x][y]
            nx, ny = x, y
            if di == 1:
                nx, ny = nx - 1, ny + 1
            else:
                nx, ny = nx + 1, ny - 1
            if idx < cnt and (nx < 0 or nx >= n or ny < 0 or ny >= m):
                if di == 1:
                    nx = x if y + 1 < m else x + 1
                    ny = y + 1 if y + 1 < m else y
                else:
                    nx = x + 1 if x + 1 < n else x
                    ny = y if x + 1 < n else y + 1
                di *= -1
            x, y = nx, ny
            idx += 1
        return ans

TypeScript 代码：

function findDiagonalOrder(g: number[][]): number[] {
    const n = g.length, m = g[0].length, cnt = n * m;
    const ans = new Array(cnt).fill(0);
    let x = 0, y = 0, dir = 1, idx = 0;
    while (idx !== cnt) {
        ans[idx++] = g[x][y];
        let nx = x, ny = y;
        if (dir === 1) {
            nx -= 1; ny += 1;
        } else {
            nx += 1; ny -= 1;
        }
        if (idx < cnt && (nx < 0 || nx >= n || ny < 0 || ny >= m)) {
            if (dir === 1) {
                nx = y + 1 < m ? x : x + 1;
                ny = y + 1 < m ? y + 1 : y;
            } else {
                nx = x + 1 < n ? x + 1 : x;
                ny = x + 1 < n ? y : y + 1;
            }
            dir *= -1;
        }
        x = nx; y = ny;
    }
    return ans;
};

• 时间复杂度：$O(n \times m)$
• 空间复杂度：$O(1)$

---

最后

这是我们「刷穿 LeetCode」系列文章的第 No.498 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。