leetcode-对角线遍历II

 题目是LeetCode第186场周赛的第三题，链接：对角线遍历 II。具体描述见原题。

 这道题的关键在于找到同一对角线上元素的关系，记一个元素的横纵坐标分别为i和j，则这里就是i+j相同的元素就处于同一对角线，而且j更大的会处于遍历顺序更后的位置。然后为了保证各条对角线的顺序（i+j小的在前），这里可以用TreeMap保存{i+j->对角线元素}的一个映射。则时间复杂度为$O(n)$，空间复杂度为$O(n)$。

 JAVA版代码如下：

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        Map<Integer, List<Integer>> idx2list = new TreeMap<>();
        int count = 0;
        int i = 0;
        for (List<Integer> num : nums) {
            int j = 0;
            for (int n : num) {
                if (idx2list.containsKey(i + j)) {
                    idx2list.get(i + j).add(n);
                }
                else {
                    List<Integer> newList = new ArrayList<>();
                    newList.add(n);
                    idx2list.put(i + j, newList);
                }
                ++j;
                ++count;
            }
            ++i;
        }
        int[] result = new int[count];
        int idx = 0;
        for (List<Integer> list : idx2list.values()) {
            for (i = list.size() - 1; i >= 0; --i) {
                result[idx++] = list.get(i);
            }
        }
        return result;
    }
}

 提交结果如下:

 为了加速，可以用数组代替TreeMap如下（本质就是空间换时间）。时间复杂度还是$O(n)$，空间复杂度为$O(n+C)$，其中$C$为$n$的最大值（这里是100000）。

 JAVA版代码如下：

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int MaxLen = 100_000;
        List<Integer>[] idx2list = new ArrayList[MaxLen];
        int count = 0;
        for (int i = nums.size() - 1; i >= 0; --i) {
            List<Integer> item = nums.get(i);
            int size = item.size();
            for (int j = 0; j < size; ++j) {
                if (idx2list[i + j] != null) {
                    idx2list[i + j].add(item.get(j));
                }
                else {
                    List<Integer> newList = new ArrayList<>();
                    newList.add(item.get(j));
                    idx2list[i + j] = newList;
                }
                ++count;
            }
        }
        int[] result = new int[count];
        int idx = 0;
        for (List<Integer> item : idx2list) {
            if (item == null) {
                break;
            }
            for (int value : item) {
                result[idx++] = value;
            }
        }
        return result;
    }
}

 提交结果如下:

 Python版代码如下：

class Solution:
    def maxScore(self, cardPoints: List[int], k: int) -> int:
        N = len(cardPoints)
        result = sum(cardPoints[0 : k])
        curSum = result
        for i in range(k):
            curSum += cardPoints[N - 1 - i] - cardPoints[k - 1 - i]
            result = max(result, curSum)
        return result

 提交结果如下: