[Leetcode] Problem 1424 - Diagonal Traverse II

Given a list of lists of integers, nums, return all elements of nums in diagonal order as shown in the below images.

Example

No.1

wmwAnH.png

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

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

No.2

wmwEBd.png

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]

No.3

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

Output: [1,4,2,5,3,8,6,9,7,10,11]

No.4

Input: nums = [[1,2,3,4,5,6]]

Output: [1,2,3,4,5,6]

Constraints

  • 1 <= nums.length <= 10^5
  • 1 <= nums[i].length <= 10^5
  • 1 <= nums[i][j] <= 10^9
  • There at most 10^5 elements in nums.

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
public int[] findDiagonalOrder(List<List<Integer>> nums) {
Map<Integer, List<Integer>> map = new HashMap<>();
int count = 0;
int idx = 0;

for (int i = 0; i < nums.size(); i++) {
List<Integer> num = nums.get(i);

for (int j = 0; j < num.size(); j++) {
map.putIfAbsent(i + j, new ArrayList<>());
map.get(i + j).add(num.get(j));
count++;
}
}

int[] result = new int[count];

for (int i = 0; i < map.size(); i++) {
List<Integer> num = map.get(i);

for (int j = num.size() - 1; j >= 0; j--) {
result[idx] = num.get(j);
idx++;
}
}

return result;
}