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.
This problem can be viewed as a graph traversal problem where each cell (i, j) is a node. The specific diagonal traversal order can be achieved using a Breadth-First Search (BFS) algorithm. We start at the top-left corner (0, 0) and explore the grid level by level, where each "level" corresponds to a diagonal.
Algorithm
- Initialize a
Queueof integer arrays (to store[row, col]coordinates) and add the starting cell(0, 0). - Initialize an empty
List<Integer>to store the traversal result. - Loop as long as the queue is not empty:
- Dequeue a cell
(r, c). - Add the value
nums.get(r).get(c)to the result list. - To explore subsequent cells, add them to the queue based on a specific rule to ensure each cell is visited exactly once and in the correct order:
- If the current cell is in the first column (
c == 0) and there is a row below it (r + 1 < nums.size()), enqueue the cell below:(r + 1, c). - If there is a cell to the right in the same row (
c + 1 < nums.get(r).size()), enqueue that cell:(r, c + 1).
- If the current cell is in the first column (
- Dequeue a cell
- After the loop finishes, convert the result list into an integer array and return it.
Walkthrough
A BFS approach uses a queue to manage the order of cells to visit. We begin by adding the starting cell (0, 0) to the queue.
The main challenge is to define the neighbors of a cell (r, c) such that we traverse diagonally and visit each cell exactly once. A clever way to achieve this is by defining the expansion rule as follows:
- From any cell
(r, c), we can always move to the right neighbor(r, c + 1)if it exists. This continues the traversal along the next diagonal. - From a cell
(r, c), we only consider moving to the cell below,(r + 1, c), if we are in the first column (c == 0). This move is responsible for starting the traversal of a new diagonal from the next row.
This set of rules guarantees that every cell is enqueued exactly once, eliminating the need for a visited set. The FIFO (First-In, First-Out) nature of the queue ensures that we process all elements of one diagonal before moving to the next, thus producing the desired output order.
class Solution { public int[] findDiagonalOrder(List<List<Integer>> nums) { Queue<int[]> queue = new LinkedList<>(); queue.offer(new int[]{0, 0}); // Start with cell (0, 0) List<Integer> resultList = new ArrayList<>(); while (!queue.isEmpty()) { int[] cell = queue.poll(); int r = cell[0]; int c = cell[1]; resultList.add(nums.get(r).get(c)); // If we are in the first column, we can start a new diagonal from the row below. if (c == 0 && r + 1 < nums.size()) { queue.offer(new int[]{r + 1, c}); } // Always try to move to the right in the current row. if (c + 1 < nums.get(r).size()) { queue.offer(new int[]{r, c + 1}); } } // Convert the list to an array for the final output. int[] result = new int[resultList.size()]; for (int i = 0; i < resultList.size(); i++) { result[i] = resultList.get(i); } return result; }}Complexity
Time
O(N), where N is the total number of elements. Each cell `(r, c)` is enqueued and dequeued exactly once.
Space
O(sqrt(N)) auxiliary space on average. The space is determined by the maximum size of the queue, which corresponds to the number of elements on the longest diagonal. For a square-like matrix, this is O(sqrt(N)). In the worst-case (e.g., a tall and skinny matrix), it can be O(N).
Trade-offs
Pros
It's a single-pass solution that builds the result directly.
Generally more space-efficient than the HashMap approach in terms of auxiliary space. The space required for the queue depends on the maximum width of a diagonal, which is often much smaller than the total number of elements N.
Cons
The logic for adding neighbors to the queue is subtle and must be implemented carefully to avoid errors like duplicate visits or incorrect order.
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; }}class Solution {public int[] findDiagonalOrder(List<List<Integer>> nums) { List<int[]> arr = new ArrayList<>(); for (int i = 0; i < nums.size(); ++i) { for (int j = 0; j < nums.get(i).size(); ++j) { arr.add(new int[]{i + j, j, nums.get(i).get(j)}); } } arr.sort((a, b)->a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]); int[] ans = new int[arr.size()]; for (int i = 0; i < arr.size(); ++i) { ans[i] = arr.get(i)[2]; } return ans; }}class Solution: def findDiagonalOrder(self, nums: List[List[int]]) -> List[int]: arr = [] for i, row in enumerate(nums): for j, v in enumerate(row): arr . append((i + j, j, v)) arr . sort() return [v[2] for v in arr]class Solution {public: vector<int> findDiagonalOrder(vector<vector<int>> &nums) { vector<tuple<int, int, int>> arr; for (int i = 0; i < nums.size(); ++i) { for (int j = 0; j < nums[i].size(); ++j) { arr.push_back({i + j, j, nums[i][j]}); } } sort(arr.begin(), arr.end()); vector<int> ans; for (auto &e : arr) { ans.push_back(get<2>(e)); } 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.