LeetCode 1424. Diagonal Traverse II

原题链接在这里：https://leetcode.com/problems/diagonal-traverse-ii/description/

题目：

Given a 2D integer array nums, return all elements of nums in diagonal order as shown in the below images.

Example 1:

Input: nums = [[1,2,3],[4,5,6],[7,8,9]]
Output: [1,4,2,7,5,3,8,6,9]

Example 2:

Input: nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
Output: [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]

Constraints:

• 1 <= nums.length <= 105
• 1 <= nums[i].length <= 105
• 1 <= sum(nums[i].length) <= 105
• 1 <= nums[i][j] <= 105

题解：

In the matrix, when index sum i + j is equal, then elements with equal sum index are on the same printing line.

Have a map to maintain the index sum and its elements.
Print from small to large index sum, for each list, print from the end.

Time Complexity: O(m * n). m = nums.size(). n is average length of list in the nums.

Space: O(m * n).

AC Java:

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int count = 0;
        int maxKey = 0;
        HashMap<Integer, ArrayList<Integer>> hm = new HashMap<>();
        for(int i = 0; i < nums.size(); i++){
            for(int j = 0; j < nums.get(i).size(); j++){
                int sum = i + j;
                hm.computeIfAbsent(sum, k -> new ArrayList<Integer>()).add(nums.get(i).get(j));
                count++;
                maxKey = Math.max(maxKey, sum);
            }
        }

        int[] res = new int[count];
        int ind = 0;
        for(int i = 0; i <= maxKey; i++){
            if(hm.containsKey(i)){
                List<Integer> ls = hm.get(i);
                for(int j = ls.size() - 1; j >=0; j--){
                    res[ind++] = ls.get(j);
                }
            }
        }

        return res;
    }
}

类似Diagonal Traverse.