Leetcode 1424. Diagonal Traverse II

Problem Explanation

In this problem, we have a jagged array (an array of arrays where each array can have different length). The goal is to retrieve elements diagonally and put them in a new array one by one.

The order of access in the 2D jagged array (nums in this case) is obtained in reverse upper-left to lower-right diagonals order. For example,

Consider the 2D jagged array :

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

In this 2D array, the sequence of access would be [1,4,2,7,5,3,8,6,9]

Here, the number 1(first of first array) is accessed first, then we move diagonally to the bottom right, i.e., to number 4(first of second array). Then we move to top right, i.e., number 2(second of first array), then again diagonally to bottom right, i.e., 7(first of third array) and so on.

Solution

Python

The Python solution works as follows:

1
2python
3class Solution:
4    def findDiagonalOrder(self, nums: List[List[int]]) -> List[int]:
5        ans = [] # Array to store the final output
6            
7        keyToNums = {} #Dictionary to store the elements according to their respective keys
8        maxKey = 0 # To keep track of the maximum key
9
10        #Iterate over each array in the nums
11        for i in range(len(nums)):
12            # Iterate over each element in the current array
13            for j in range(len(nums[i])):
14                key = i + j # Calculate the key for the element
15                if key in keyToNums: # If the key is already present in the dictionary, append the current element in the value of that key
16                    keyToNums[key].append(nums[i][j])
17                else:
18                    keyToNums[key] = [nums[i][j]] #Else, assign a new array to the key with current element as the first element of the array
19
20                maxKey = max(maxKey, key) # Update the maximum key if necessary
21        
22        #Iterate over each key till the maximum key
23        for i in range(maxKey+1):
24            # For each key, the elements are to be returned in reverse order of their appearance
25            ans += keyToNums[i][::-1]
26            
27        return ans

JavaScript

1
2javascript
3var findDiagonalOrder = function (nums) {
4    const ouputArray = [];
5    const map = new Map();
6    let maxKey = 0;
7    for (let i = 0; i < nums.length; i++) {
8        for (let j = 0; j < nums[i].length; j++) {
9            const key = i + j;
10            if (!map.has(key)) {
11                map.set(key, [nums[i][j]]);
12            } else {
13                map.get(key).push(nums[i][j]);
14            }
15            maxKey = Math.max(maxKey, key);
16        }
17    }
18    for (let i = 0; i <= maxKey; i++) {
19        ouputArray.push(...map.get(i).reverse());
20    }
21
22    return ouputArray;
23};

Java

1
2java
3class Solution {
4    public int[] findDiagonalOrder(List<List<Integer>> nums) {
5        int n = nums.size(), maxLen = 0, total = 0;
6        HashMap<Integer, List<Integer>> map = new HashMap<>();
7        for (int r = n - 1; r >= 0; --r) {
8            for (int c = nums.get(r).size() - 1; c >= 0; --c) {
9                map.putIfAbsent(r + c, new ArrayList<>());
10                map.get(r + c).add(nums.get(r).get(c));
11                total++;
12            }
13            maxLen = Math.max(maxLen, nums.get(r).size());
14        }
15        int[] res = new int[total];
16        int idx = 0;
17        for (int d = 0; d <= n + maxLen; ++d) {
18            List<Integer> list = map.get(d);
19            if (list == null) continue;
20            for (int i = list.size() - 1; i >= 0; --i)
21                res[idx++] = list.get(i);
22        }
23        return res;
24    }
25}

C++

1
2c++
3class Solution {
4public:
5    vector<int> findDiagonalOrder(vector<vector<int>>& nums) {
6        unordered_map<int,vector<int>>m;
7        int maxr=0;
8        for(int i=0;i<nums.size();i++)
9        {
10            for(int j=0;j<nums[i].size();j++)
11            {
12                m[i+j].push_back(nums[i][j]);
13                maxr=max(maxr,(int)m[i+j].size());
14            }
15        }
16        vector<int>result;
17        for(int i=0;i<=maxr+nums.size();i++)
18        {
19            for(int j=m[i].size()-1;j>=0;j--)
20            {
21                result.push_back(m[i][j]);
22            }
23        }
24        return result;
25    }
26};

C#

1
2csharp
3public class Solution {
4    public IList<int> FindDiagonalOrder(IList<IList<int>> nums) {
5        // we need to use a linked list since it has O(1) time complexity to remove last
6        var map = new Dictionary<int, LinkedList<int>>();
7        int size = 0, maxKey = 0;
8        // add elements to map and switch the row with column for each cell to remember diagonals
9        // record the count of each sum of row and column
10        foreach (var lst in nums.Select((Value, Key) => new { Key, Value })) { // use to get row number
11            int r = lst.Key, n = lst.Value.Count;
12            size += n;
13            maxKey = Math.Max(maxKey, r + n - 1); // record maximum diagonal index
14            for (int c = 0; c < n; ++c) {
15                if (map.ContainsKey(r + c)) { // row + column number is the index of diagonal
16                    map[r + c].AddLast(lst.Value[c]);
17                } else {
18                    map[r + c] = new LinkedList<int>(new List<int>() {lst.Value[c]});
19                }
20            }
21        }
22        
23        var result = new List<int>();
24        for (int i = 0; i <= maxKey; ++i) {
25            // append list from the back to mimic the traversal order
26            if (map.ContainsKey(i)) {
27                result.AddRange(map[i].Reverse());
28            }
29        }
30        return result;
31    }
32}

Discussion

As you can see, all of these solutions use essentially the same logic. They iterate over each item in the jagged array, map each item to a new array or list based on its diagonal, and then append each element to a separate list, array or vector in reverse order. This approach works well because it follows the problem's requirements closely, ensuring that we access all elements in the correct order.

In Python, JavaScript, and Java versions, a Dictionary or a Map acts as the main data structure for mapping elements according to their diagonal indices. Dictionary in Python, Map in JavaScript and HashMap in Java have constant time complexity O(1) for inserting and accessing elements which ensures efficiency during operations.

In C++, an unordered_map is used as it acts similarly as Map in JavaScript, Dictionary in Python and HashMap in Java.

For C#, a Dictionary with LinkedLists of Integers is used. LinkedList is chosen over a List here because it provides O(1) time complexity while removing the last element. It avoids shifting the remaining elements when an element is removed from the end of the list, leading to higher efficiency.

The Python solution uses a list comprehension to reverse the appearance order, while the JavaScript and Java versions use traditional loops. The C++ version also uses a traditional loop but it pushes back the elements directly into the resultant vector.

In summary, these solutions effectively traverse a jagged array in a diagonal order diligently following the access sequence as required by the problem. The results are then compiled and returned finally indicating the end of diagonal traversal.

It's also important to note that these solutions handle both standard 2D arrays and jagged arrays, ensuring their generalizability.


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫