1424. 对角线遍历 II：数组问题，优先队列或排序？但我们真的需要优先队列或排序吗？.md

题目简述：

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

题目链接：1424. 对角线遍历 II

暴力 BFS（超时）

思路

最初我的想法是直接模拟，按矩阵对角线的顺序遍历。如果矩阵上的某个点实际上不存在，就跳过，然后继续考察下一个可能的点。

但这样子实现起来算法的时间复杂度就不是 $O(n)$ 了，而是 $O(size\times\max n_k)$，其中 $n$ 是列表中数字的总数，$size$ 是整数列表的数量，$n_k$ 是第 $k$ 个整数列表的长度。

下面就是按上述思路实现的一个算法，由于 $size\times\max n_k\geqslant n$，且在一些极端的测试用例下 $size\times\max n_k$ 远大于 $n$，所以该算法效率并不高。实际上这份代码在力扣平台上执行超时了，说明这个实现亟需改进。

代码

算法的时间复杂度为 $O(size\times\max size_k)$，空间复杂度为 $O(n)$。

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int len = 0;
        int m = nums.size();
        int n = 0;
        for (List<Integer> arr : nums) {
            n = Math.max(n, arr.size());
            len += arr.size();
        }

        int[] res = new int[len];
        int i = 0;
        for (int level = 0; level < m + n - 1; level++) {
            int row = Math.min(level, m - 1);
            int col = level - row;
            while (row >= 0 && col < n) {
                if (col < nums.get(row).size()) res[i++] = nums.get(row).get(col);
                row--;
                col++;
            }
        }

        return res;
    }
}

优先队列 & 排序

思路

查阅题目标签后，我发现本问题可以利用优先队列这一数据结构解决。于是稍加思索，我就想出来了这一解法。

考虑按对角线的序（也可以解释为 BFS 的层数）$row+col$ 对各元素分桶，也就是说，当且仅当 $row_1+col_1=row_2+col_2$ 时，视 $nums[row_1][col_1]$ 与 $nums[row_2][col_2]$ 为同一分类下的元素，归结于同一个桶内。然后，对同一个分类桶内的元素按行号或列号维护顺序，这样就可以保证：

• 粗颗粒度上，不同桶之间按对角线的序（$row+col$，现在可以视为桶编号）决定遍历的相对顺序；
• 细颗粒度上，同一桶内的各元素按行号或列号决定遍历的相对顺序。

二者结合，我们便只需要对 nums 进行一次完整的遍历，就可以得到题目所要求的遍历顺序。

在同一桶内，可以用优先队列按行号或列号维护顺序，也可以在遍历结束后分别对各桶内元素按行号或列号显式地进行排序。

代码

这里考虑优先队列的实现，排序的方案是同理的。算法的时间复杂度为 $O(n\log n)$，空间复杂度为 $O(n)$。

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int n = 0;
        Map<Integer, PriorityQueue<Pair>> map = new HashMap<>();
        for (int row = 0; row < nums.size(); row++) {
            List<Integer> line = nums.get(row);
            n += line.size();
            for (int col = 0; col < line.size(); col++) {
                map.computeIfAbsent(row + col, k -> new PriorityQueue<>()).add(new Pair(col, line.get(col)));
            }
        }

        int[] res = new int[n];
        int q = 0;
        for (int i = 0; ; i++) {
            PriorityQueue<Pair> pq = map.get(i);
            if (pq == null) break;
            while (!pq.isEmpty()) res[q++] = pq.poll().val();
        }
        return res;
    }

    record Pair(int col, int val) implements Comparable<Pair> {
        @Override
        public int compareTo(Pair other) {
            return Integer.compare(this.col, other.col);
        }
    }
}

我们真的需要排序吗？

思路

其实在优先队列代码的基础上我们还可以实现时间复杂度为 $O(n)$ 的算法。

这里有一个 trick：同一对角线内，nums 顺序遍历的序（即输入的序）与对角线遍历的序是相反的。读者不妨试着证明一下。

所以，只要我们使用某种方法让遍历反向，无论是头插还是使用栈，甚至先尾插再 reverse，都能避免排序带来的 $O(n\log n)$ 时间开销。

代码

算法的时间复杂度为 $O(n)$，如果不考虑存储结果所需要的空间开销，则空间复杂度可以认为是 $O(1)$ 的。

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int n = 0;
        Map<Integer, Deque<Integer>> map = new HashMap<>();
        for (int row = 0; row < nums.size(); row++) {
            List<Integer> line = nums.get(row);
            n += line.size();
            for (int col = 0; col < line.size(); col++) {
                map.computeIfAbsent(row + col, k -> new ArrayDeque<>()).addFirst(line.get(col));
            }
        }

        Deque<Integer> res = map.get(0);
        for (int i = 1; ; i++) {
            Deque<Integer> queue = map.get(i);
            if (queue == null) break;
            res.addAll(queue);
        }
        return res.stream().mapToInt(Integer::intValue).toArray();
    }
}