Medium
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
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
public class Solution {
public int[] findDiagonalOrder(List<List<Integer>> nums) {
List<Integer> ans = new ArrayList<>();
ArrayDeque<Iterator<Integer>> queue = new ArrayDeque<>();
int pos = 0;
do {
if (pos < nums.size()) {
queue.offerFirst(nums.get(pos).iterator());
}
int sz = queue.size();
while (--sz >= 0) {
Iterator<Integer> cur = queue.poll();
ans.add(Objects.requireNonNull(cur).next());
if (cur.hasNext()) {
queue.offer(cur);
}
}
pos++;
} while (!queue.isEmpty() || pos < nums.size());
return ans.stream().mapToInt(o -> o).toArray();
}
}