Diagonal Traverse II
MedPrompt
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 <= 1051 <= nums[i].length <= 1051 <= sum(nums[i].length) <= 1051 <= nums[i][j] <= 105
Approaches
2 approaches with complexity analysis and trade-offs.
This approach leverages the property that all elements on the same diagonal share the same sum of their row and column indices (i + j). We can iterate through all the elements of the input nums, calculate this sum for each element, and use a HashMap to group elements by their corresponding sum. This effectively collects all elements for each diagonal into separate lists.
Algorithm
- Create a
HashMap<Integer, List<Integer>>to group elements by their diagonal sum, which isrow_index + column_index. - To ensure the correct order within each diagonal (bottom-up), iterate through the input
numsfrom the last row to the first row (ifromnums.size() - 1down to0). - For each element
nums[i][j], calculate its diagonal keykey = i + j. - Add the element to the list associated with this
keyin the HashMap. If the key is new, create a new list first. - While iterating, keep track of the total number of elements
Nand the maximum diagonal keymaxKeyencountered. - After populating the map, create a result array of size
N. - Iterate from
key = 0tomaxKey. - For each key, retrieve the list of diagonal elements from the map and append them to the result array.
- Return the final result array.
Walkthrough
The main idea is to use a HashMap<Integer, List<Integer>> where each key represents a diagonal (via the sum i + j) and the value is a list of all numbers on that diagonal.
To achieve the required output order (diagonals from top-left to bottom-right, and within each diagonal, from bottom-left to top-right), we need to be careful about how we populate these lists. A simple and effective way is to iterate through the input nums matrix starting from the last row and moving upwards. By doing this, for any given diagonal sum k = i + j, we will process elements with a larger row index i first. When we add these elements to the end of the list for key k, they are naturally placed in the correct bottom-to-top order.
After iterating through all the elements and populating the map, we can construct the final result. We iterate through the diagonal keys from 0 up to the maximum key found. For each key, we append the list of elements from our map to the final result array.
class Solution { public int[] findDiagonalOrder(List<List<Integer>> nums) { Map<Integer, List<Integer>> groups = new HashMap<>(); int n = 0; int maxKey = 0; // Iterate from bottom-to-top, left-to-right to ensure correct order within diagonals for (int i = nums.size() - 1; i >= 0; i--) { n += nums.get(i).size(); for (int j = 0; j < nums.get(i).size(); j++) { int key = i + j; groups.putIfAbsent(key, new ArrayList<>()); groups.get(key).add(nums.get(i).get(j)); maxKey = Math.max(maxKey, key); } } int[] result = new int[n]; int index = 0; // Reconstruct the result from the map for (int key = 0; key <= maxKey; key++) { List<Integer> diagonal = groups.get(key); if (diagonal != null) { for (int val : diagonal) { result[index++] = val; } } } return result; }}Complexity
Time
O(N), where N is the total number of elements in `nums`. We perform a single pass through all N elements to populate the map, and another pass through the N elements to build the result array.
Space
O(N), where N is the total number of elements in `nums`. The HashMap stores all N elements, and the result array also stores N elements.
Trade-offs
Pros
The logic is straightforward and directly models the definition of a diagonal.
Relatively easy to implement correctly.
Cons
Requires O(N) auxiliary space to store all elements in the HashMap, which can be substantial for large inputs.
Effectively requires two passes over the data: one to populate the map and another to build the result array from the map.
Solutions
Solution
public class Solution { public int[] FindDiagonalOrder(IList < IList < int >> nums) { List < int[] > arr = new List < int[] > (); for (int i = 0; i < nums.Count; ++i) { for (int j = 0; j < nums[i].Count; ++j) { arr.Add(new int[] { i + j, j, nums[i][j] }); } } arr.Sort((a, b) => a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]); int[] ans = new int[arr.Count]; for (int i = 0; i < arr.Count; ++i) { ans[i] = arr[i][2]; } return ans; }}Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.