1424. Diagonal Traverse II

JIAKAOBO

LeetCode

感谢赞助!

• ㊗️
• 大家
• offer
• 多多！

Problem

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

Code

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int count = 0, maxKey = 0;
        Map<Integer, List<Integer>> map = new HashMap<>();

        for (int row = nums.size() - 1; row >= 0; row--) {
            for (int col = 0; col < nums.get(row).size(); col++) {
                int key = row + col;
                if (!map.containsKey(key)) {
                    map.put(key, new ArrayList<>());
                }
                map.get(key).add(nums.get(row).get(col));
                maxKey = Math.max(maxKey, row + col);
                count++;
            }
        }

        int[] ans = new int[count];
        int index = 0;

        for (int key = 0; key <= maxKey; key++) {
            if (!map.containsKey(key))
                continue;
            for (int v : map.get(key))
                ans[index++] = v;
        }

        return ans;
    }
}