排序题目：对角线遍历 II

题目

标题和出处

标题：对角线遍历 II

出处：1424. 对角线遍历 II

难度

6 级

题目描述

要求

给定一个二维整数数组 $\text{nums}$，将 $\text{nums}$ 中的所有元素按照如下图所示的对角线顺序返回。

示例

示例 1：

输入：$\text{nums = [[1,2,3],[4,5,6],[7,8,9]]}$
输出：$\text{[1,4,2,7,5,3,8,6,9]}$

示例 2：

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

数据范围

• $\text{1}\le \text{nums.length}\le {\text{10}}^{\text{5}}$
• $\text{1}\le \text{nums[i].length}\le {\text{10}}^{\text{5}}$
• $\text{1}\le \text{sum(nums[i].length)}\le {\text{10}}^{\text{5}}$
• $\text{1}\le \text{nums[i][j]}\le {\text{10}}^{\text{5}}$

解法

思路和算法

这道题给定的二维数组 $\text{nums}$ 最多有 $10^5$ 行和 $10^5$ 列，如果直接遍历二维数组中的每个位置，则最多需要遍历 $10^{10}$ 个位置，会超出时间限制，因此需要考虑时间复杂度更低的解法。

由于 $\text{nums}$ 不是完整的矩阵，可能有部分位置并没有元素，元素总数最多是 $10^5$，因此只需要遍历有元素的位置，并将元素排序。

根据对角线遍历的规则，每条对角线的方向是从左下到右上，依次遍历每条对角线，同一条对角线上的元素顺序是列下标升序顺序（或者行下标降序顺序）。

由于同一条对角线上的每个位置的行下标与列下标之和是定值，因此可以使用行下标与列下标之和表示对角线编号，对角线编号用于区分不同的对角线。

遍历 $\text{nums}$ 并使用列表存储每个元素的信息，每个元素需要记录三个信息：元素所在的行下标与列下标之和、元素所在的列下标、元素值，其中元素所在的行下标与列下标之和即为对角线编号。

遍历之后对列表排序，排序的依据如下：

1. 如果两个元素的对角线编号不同，则根据对角线编号升序排序；

2. 如果两个元素的对角线编号相同，则根据列下标升序排序。

遍历排序后的列表，对于列表中的每个元素，将元素值填入结果数组中。

代码

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        List<int[]> list = new ArrayList<int[]>();
        int rows = nums.size();
        for (int i = 0; i < rows; i++) {
            List<Integer> rowList = nums.get(i);
            int cols = rowList.size();
            for (int j = 0; j < cols; j++) {
                int num = rowList.get(j);
                list.add(new int[]{i + j, j, num});
            }
        }
        Collections.sort(list, (a, b) -> {
            if (a[0] != b[0]) {
                return a[0] - b[0];
            } else {
                return a[1] - b[1];
            }
        });
        int size = list.size();
        int[] order = new int[size];
        for (int i = 0; i < size; i++) {
            order[i] = list.get(i)[2];
        }
        return order;
    }
}

复杂度分析

• 时间复杂度：$O(n\log n)$，其中 $n$ 是二维数组 $\text{nums}$ 中的元素个数。排序需要 $O(n\log n)$ 的时间，每次遍历都需要 $O(n)$ 的时间，因此时间复杂度是 $O(n\log n)$。

• 空间复杂度：$O(n)$，其中 $n$ 是二维数组 $\text{nums}$ 中的元素个数。列表需要 $O(n)$ 的空间。