算法题003：对角线遍历 II

发表于 2023-10-29 分类于算法题本文字数： 1.1k 阅读时长 ≈ 4 分钟

题目

题目来源：1424. 对角线遍历 II

给你一个列表 nums ，里面每一个元素都是一个整数列表。请你依照下面各图的规则，按顺序返回 nums 中对角线上的整数。

示例 1

输入：nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]] 输出：[1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]

示例 2

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

示例 3

输入：nums = [[1,2,3],[4],[5,6,7],[8],[9,10,11]] 输出：[1,4,2,5,3,8,6,9,7,10,11]

数据范围

1 <= nums.length <= 10^5
1 <= nums[i].length <= 10^5
1 <= nums[i][j] <= 10^9
nums 中最多有 10^5 个数字。

思路

对于自然数 $k \in \mathbb{N}$，第 $k + 1$ 条对角线上的元素坐标为 $\set{(i,j) | i + j = k}$。例如，当 $k = 2$ 时，第 $3$ 条对角线的元素坐标是 $\set{(0,2),(1,1),(2,0)}$

如果 nums 是个矩阵，那这题是一道简单模拟题，按照对角线遍历即可，时间复杂度为 $O(mn)$。

但是，nums 的行数 $m \le 10^5$，列数 $n \le 10^5$，$O(mn)$ 的时间复杂度会超时。因此，我们不应该把 nums 当作一个矩阵来遍历，否则会访问大量的空格子。

观察示例1：

观察1️⃣：对于任意的 $k \in \mathbb{N}$，第 $k + 1$ 条对角线的元素遍历顺序是从左下到右上，即 $(k, 0), (k - 1,1),\cdots,(0,k)$。因此，如果我们要获得第 $k + 1$ 条对角线上的元素，我们应该倒序遍历 nums 中每一行的对应元素。

观察2️⃣：易得，对于nums中的每一行 nums[j]，其中的元素都属于不同的对角线，因此是相互独立的。

观察3️⃣：倒序遍历 nums 中每一行时：

对于最后一行的元素 $12,13,14,15,16$，它们都是自己对角线上的第一个元素，尽管后 4 个元素所属的对角线在第一列上不存在元素。
对于倒数第二行的元素 $9,10,11$，它们都是所属对角线的下一个应该遍历的元素
以此类推……

我们的目标是要收集 nums 中的所有元素，收集方式是按照对角线元素顺序收集。

因此，我们可以先为每个对角线创建一个 list，将对应元素放入该 list 中。然后再正序遍历每个对角线 list，获得最终的结果。如果这样做，我们需要首先创建 m + n - 1 个 list。

int m = nums.size();
// n 为最大列数
int n = nums.stream().mapToInt(List::size).max().getAsInt();
// ans 中存放 m + n - 1 个 list，代表各个对角线中遍历到的元素
List<List<Integer>> ans = new ArrayList<>();
IntStream.range(0, m + n - 1).forEach(i -> ans.add(new ArrayList<>()));

// 核心：如何遍历？？？

// 正序收集 ans 中每行的元素
return ans.stream().flatMap(l -> l.stream()).mapToInt(i -> i).toArray();

那么如何遍历呢？为了不超时，肯定是按照行遍历。

根据『观察1️⃣』，倒序遍历 nums ；
根据『观察2️⃣』，对于每一行，每遍历到一个元素 nums[i][j]，就把它放入第 $i+j + 1$ 个 list 中
根据『观察3️⃣』，上述元素 nums[i][j] 应该添加到对应 list 最后。因为前 $n - i - 1$ 行（$n -1,n - 2 ,\cdots, i + 1$）中该对角线上的元素已经遍历过了，并且它们已经按照同样的方法、有序存放在该 list 中了—— 使用递归或者数学归纳法证明

代码是简单的：

// 倒序遍历 nums
for (int i = m - 1; i >= 0; i--) {
    List<Integer> row = nums.get(i);
    // 每遍历到一个元素 nums[i][j]，就把它放入第 i + j + 1 个 list 中
    for (int j = 0; j < row.size(); j++) {
        ans.get(i + j).add(row.get(j));
    }
}

时间复杂度是 $O(m + n + \text{(nums 元素个数)})$

空间复杂度是 $O(m+n)$

代码

Java

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int m = nums.size();
        // n 为最大列数
        int n = nums.stream().mapToInt(List::size).max().getAsInt();
        // ans 中存放 m + n - 1 个 list，代表各个对角线中遍历到的元素
        List<List<Integer>> ans = new ArrayList<>();
        IntStream.range(0, m + n - 1).forEach(i -> ans.add(new ArrayList<>()));

        // 倒序遍历 nums
        for (int i = m - 1; i >= 0; i--) {
            List<Integer> row = nums.get(i);
            // 每遍历到一个元素 nums[i][j]，就把它放入第 i + j + 1 个 list 中
            for (int j = 0; j < row.size(); j++) {
                ans.get(i + j).add(row.get(j));
            }
        }
        // 正序收集 ans 中每行的元素
        return ans.stream().flatMap(l -> l.stream()).mapToInt(i -> i).toArray();
    }
}