1424. Diagonal Traverse II


Problem Description

Given a 2-dimensional grid nums, where each cell contains an integer, our task is to return a list of all the elements of nums but in a specific order—diagonal order. In this ordering, elements from the same diagonal (elements with the same sum of their row and column indices) should be grouped together. Additionally, within each diagonal group, elements should be sorted based on their column indices. This creates a zigzag pattern if we visualize it, as we traverse diagonals starting from the top row, moving down-right on the grid, and then up to the next diagonal start point.

The challenge is to implement this in a way that respects the diagonal grouping and within-group ordering, and to do so as efficiently as possible.

Intuition

The intuition behind the solution lies in recognizing that elements that are on the same diagonal share a certain property: the sum of their row and column indices is constant. With this realization, we can iterate through every element in the nums grid, calculate the sum of its indices, and use this sum as a way to group elements into their respective diagonals.

To maintain the within-diagonal order (where elements with lower column indices come first), we keep track of the column index alongside the diagonal index sum. This means that for each element, we store a triple containing the sum (which represents the diagonal), the column index, and the value itself.

Once we have gone through all the elements and collected this information, we sort this list of triples. The sort will naturally group elements from the same diagonal together (because of the sums), and within each group, it will order elements by their column indices.

Finally, we extract the third element of each triple (which is the value from the original grid) to get our result in the correct diagonal order. The fact that the Python sort function sorts tuples lexicographically, first by the first element, then by the second, etc., makes it particularly well-suited for this task.

Learn more about Sorting and Heap (Priority Queue) patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece:

Which type of traversal does breadth first search do?

Solution Approach

The implementation of the solution makes efficient use of Python's list and sorting capabilities.

Here's a step-by-step explanation of the approach:

  1. Initialize an empty list arr: This list is used to store the triples (sum of indices, column index, value) for each element in the nums array.

  2. Iterate over the grid: We use a nested loop to go through each element of the nums grid. The outer loop (for i, row in enumerate(nums)) iterates over the rows, and the inner loop (for j, v in enumerate(row)) goes over each element in a row. Here, i is the row index, j is the column index, and v is the value at the nums[i][j] position.

  3. Compute the diagonal index and store triples: For each grid element nums[i][j], we compute the sum of the indices (i + j) which uniquely identifies the diagonal to which the element belongs. We then create a triple (i + j, j, v) and append it to our arr list. The triple contains the diagonal index, the column index, and the value itself, respectively.

  4. Sort the list of triples: We call arr.sort() to sort the list of triples. The sort method uses the natural lexicographical order for tuples, so it will first sort them by the diagonal index, and then by their column index within the same diagonal. This operation arranges the elements in exactly the order required for the final result.

  5. Extract the results: The last line is a list comprehension [v[2] for v in arr] that goes through the sorted list of triples and extracts the third element from each triple (the value). This creates a new list of integers which is our final, diagonally ordered list of elements.

The key algorithms and data structures used in this approach are:

  • Nested Loops: To iterate over a 2D array.
  • List Comprehension: A concise way to construct lists.
  • Enumerate Function: Provides a neat way to get both the index and the value when iterating over a list.
  • Triple (Tuple): An immutable data structure to store the related elements - sum of indices, column index, and value.
  • Sorting: A built-in Python method to sort lists in ascending order based on the first element of the tuple, and if they are the same, based on the second element, and so on.

This approach has a time complexity of O(N log N), where N is the total number of elements in the nums array, because the most expensive operation is sorting the list of triples. The spatial complexity is O(N) since we need to create a list arr that holds as many triples as there are elements in the original array.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Which of the following is a min heap?

Example Walkthrough

Let's use a small 3x3 grid example to illustrate the solution approach:

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

Following the steps in the solution approach:

  1. Initialize an empty list arr:

    At the beginning, arr is an empty list: arr = []

  2. Iterate over the grid:

    Now we start the nested loop over the grid:

    • For i = 0, row = [1, 2, 3]:
      • For j = 0, v = 1: Compute diagonal index i + j = 0, append (0, 0, 1) to arr.
      • For j = 1, v = 2: Compute diagonal index i + j = 1, append (1, 1, 2) to arr.
      • For j = 2, v = 3: Compute diagonal index i + j = 2, append (2, 2, 3) to arr.
    • For i = 1, row = [4, 5, 6]:
      • For j = 0, v = 4: Compute diagonal index i + j = 1, append (1, 0, 4) to arr.
      • For j = 1, v = 5: Compute diagonal index i + j = 2, append (2, 1, 5) to arr.
      • For j = 2, v = 6: Compute diagonal index i + j = 3, append (3, 2, 6) to arr.
    • For i = 2, row = [7, 8, 9]:
      • For j = 0, v = 7: Compute diagonal index i + j = 2, append (2, 0, 7) to arr.
      • For j = 1, v = 8: Compute diagonal index i + j = 3, append (3, 1, 8) to arr.
      • For j = 2, v = 9: Compute diagonal index i + j = 4, append (4, 2, 9) to arr.

    arr now looks like this: [(0, 0, 1), (1, 1, 2), (2, 2, 3), (1, 0, 4), (2, 1, 5), (3, 2, 6), (2, 0, 7), (3, 1, 8), (4, 2, 9)]

  3. Sort the list of triples:

    After sorting, arr should look like this: [(0, 0, 1), (1, 0, 4), (1, 1, 2), (2, 0, 7), (2, 1, 5), (2, 2, 3), (3, 1, 8), (3, 2, 6), (4, 2, 9)]

    Notice the elements are grouped by the sum of their indices (the diagonals) and sorted by their column indices within each group.

  4. Extract the results:

    By extracting the third element of each tuple, we get the diagonally ordered list of elements: [1, 4, 2, 7, 5, 3, 8, 6, 9]

That's the diagonal order traversal of the 2D nums grid using the proposed solution approach.

Not Sure What to Study? Take the 2-min Quiz:

How many ways can you arrange the three letters A, B and C?

Python Solution

1from typing import List
2
3class Solution:
4    def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
5        # Initialize an empty list to store the values along with their diagonal indices
6        diagonal_elements = []
7      
8        # Iterate through the matrix rows
9        for row_index, row in enumerate(matrix):
10            # Iterate through the elements in the current row
11            for column_index, value in enumerate(row):
12                # The sum of row and column indices gives the diagonal index.
13                # Append a tuple containing the diagonal index, column index and the value
14                diagonal_elements.append((row_index + column_index, column_index, value))
15      
16        # Sort the list of tuples based on diagonal index, then by column index (as secondary)
17        # This will order the elements first by diagonal, then from top-right to bottom-left
18        # within the same diagonal line.
19        diagonal_elements.sort()
20      
21        # Extract the values from the sorted list of tuples and return them in the correct order
22        return [element[2] for element in diagonal_elements]
23

Java Solution

1import java.util.List;
2import java.util.ArrayList;
3import java.util.Collections;
4
5class Solution {
6
7    /**
8     * Function to find the diagonal order of a given list of lists of integers.
9     * Diagonals are considered in a "zig-zag"/top-right to bottom-left manner.
10     *
11     * @param nums 2D List of integers representing a matrix.
12     * @return An array of integers representing the diagonal traversal.
13     */
14    public int[] findDiagonalOrder(List<List<Integer>> nums) {
15        // Temporary list to store elements and their diagonal indices
16        List<int[]> diagonalElements = new ArrayList<>();
17
18        // Iterate over the 2D List to process each element
19        for (int i = 0; i < nums.size(); i++) {
20            for (int j = 0; j < nums.get(i).size(); j++) {
21                // Calculate the diagonal index and store it with the element
22                // The format of the stored array is [diagonal index, column index, value]
23                diagonalElements.add(new int[] {i + j, j, nums.get(i).get(j)});
24            }
25        }
26
27        // Sort the diagonalElements based on their computed diagonal indices
28        // If two elements are on the same diagonal, sort them by their column index
29        Collections.sort(diagonalElements, (a, b) -> {
30            if (a[0] == b[0]) {
31                return a[1] - b[1];
32            } else {
33                return a[0] - b[0];
34            }
35        });
36
37        // Initialize the answer array with the correct size
38        int[] ans = new int[diagonalElements.size()];
39
40        // Extract the values from the sorted diagonalElements and populate the answer array
41        for (int i = 0; i < diagonalElements.size(); i++) {
42            ans[i] = diagonalElements.get(i)[2];
43        }
44
45        // Return the final diagonal order traversal as an array
46        return ans;
47    }
48}
49

C++ Solution

1#include <vector>
2#include <tuple>
3#include <algorithm>
4
5class Solution {
6public:
7    // Function to find and return the elements of the 2D vector in diagonal order.
8    vector<int> findDiagonalOrder(vector<vector<int>>& nums) {
9        // Create a vector to store tuples containing the diagonal number, column number, and value.
10        vector<tuple<int, int, int>> diagonalElements;
11
12        // Iterate over the 2D vector to populate diagonalElements.
13        for (int row = 0; row < nums.size(); ++row) {
14            for (int col = 0; col < nums[row].size(); ++col) {
15                // The sum of row and column indices determines the diagonal number.
16                // Store diagonal number, column number, and element.
17                diagonalElements.push_back({row + col, col, nums[row][col]});
18            }
19        }
20
21        // Sort the elements of diagonalElements. By default, tuples are sorted by their
22        // first element, then by their second element, so this will sort by the diagonal number,
23        // then by the column index.
24        sort(diagonalElements.begin(), diagonalElements.end());
25
26        // Create a vector to store the answer in the diagonal order.
27        vector<int> ans;
28        // Extract the elements from the tuples in sorted order and append them to ans.
29        for (auto& element : diagonalElements) {
30            ans.push_back(get<2>(element));
31        }
32
33        // Return the answer.
34        return ans;
35    }
36};
37

Typescript Solution

1type DiagonalElement = [number, number, number];
2
3// Function to find and return the elements of the 2D array in diagonal order.
4function findDiagonalOrder(nums: number[][]): number[] {
5    // Create an array to store tuples containing the diagonal number, column number, and value.
6    let diagonalElements: DiagonalElement[] = [];
7
8    // Iterate over the 2D array to populate diagonalElements.
9    for (let row = 0; row < nums.length; ++row) {
10        for (let col = 0; col < nums[row].length; ++col) {
11            // The sum of row and column indices determines the diagonal number.
12            // Store diagonal number, column number, and element.
13            diagonalElements.push([row + col, col, nums[row][col]]);
14        }
15    }
16
17    // Sort the elements of diagonalElements. Since TypeScript requires a comparator function,
18    // it's provided here. This function sorts by the diagonal number, then by the column index.
19    diagonalElements.sort((a, b) => {
20        if (a[0] === b[0]) return a[1] - b[1]; // If diagonal numbers match, sort by column.
21        return a[0] - b[0]; // Otherwise, sort by diagonal number.
22    });
23
24    // Create an array to store the answer in the diagonal order.
25    let answer: number[] = [];
26    // Extract the elements from the tuples in sorted order and append them to answer.
27    for (let element of diagonalElements) {
28        answer.push(element[2]);
29    }
30
31    // Return the answer.
32    return answer;
33}
34
35// Example of how to use the function.
36// const matrix: number[][] = [
37//     [1, 2, 3],
38//     [4, 5, 6],
39//     [7, 8, 9]
40// ];
41// const diagonalOrder: number[] = findDiagonalOrder(matrix);
42
Fast Track Your Learning with Our Quick Skills Quiz:

What is the running time of the following code?

1int sqrt(int n) {
2  for (int guess = 1; guess * guess <= n; guess++) {
3    if (guess * guess == n) {
4      return guess;
5    }
6  }
7  return -1;
8}

Time and Space Complexity

The time complexity of the code is determined by the number of operations it performs. The given Python code iterates over all the elements in the list of lists to create a flat list of tuples (arr). The iteration has a complexity of O(N), where N is the total number of elements in the nums list of lists. Additionally, the code sorts the arr list, which has O(N log N) complexity for the Timsort algorithm used in Python's sort() method. Since the sort is the most expensive operation here, the overall time complexity is O(N log N).

The space complexity of the code considers the additional space used by the algorithm outside the input and output data. The main extra space used is for the arr list of tuples. Since each element in the input is transformed into a tuple and stored in arr, the space complexity is O(N), where N is again the total number of elements in nums.

Learn more about how to find time and space complexity quickly.


Recommended Readings


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 👨‍🏫