Problem Description

You are given a 2D integer array nums where rows can have different lengths (jagged array). Your task is to return all elements of nums in diagonal order.

The diagonal order means traversing the array along diagonals that go from top-left to bottom-right. Start with the diagonal containing nums[0][0], then move to the next diagonal containing nums[0][1] and nums[1][0], and so on. Within each diagonal, elements should be ordered from bottom-left to top-right (i.e., elements with larger row indices come first).

For example, if you have:

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

The diagonals would be:

• First diagonal: [1] (position [0][0])
• Second diagonal: [4, 2] (positions [1][0], [0][1])
• Third diagonal: [6, 5, 3] (positions [2][0], [1][1], [0][2])
• Fourth diagonal: [7] (position [2][1])
• Fifth diagonal: [8] (position [2][2])
• Sixth diagonal: [9] (position [2][3])

The final result would be: [1, 4, 2, 6, 5, 3, 7, 8, 9]

The solution leverages the observation that elements on the same diagonal have the same sum of indices i + j. By creating tuples of (i+j, j, value) for each element and sorting them, we ensure that elements are grouped by diagonal (same i+j value) and within each diagonal, elements with smaller j values (which correspond to larger i values) come first.

Intuition

To understand how to traverse diagonally, let's first identify what makes elements belong to the same diagonal. If we look at the indices of elements on the same diagonal:

• Diagonal 1: [0,0]
• Diagonal 2: [1,0], [0,1]
• Diagonal 3: [2,0], [1,1], [0,2]

We notice that for each diagonal, the sum i + j remains constant:

• Diagonal 1: 0 + 0 = 0
• Diagonal 2: 1 + 0 = 1, 0 + 1 = 1
• Diagonal 3: 2 + 0 = 2, 1 + 1 = 2, 0 + 2 = 2

This gives us our first key insight: elements with the same i + j value belong to the same diagonal.

Next, we need to determine the order within each diagonal. Looking at the required output order within a diagonal (from bottom-left to top-right), we see that elements with larger row indices come first. Since i + j is constant within a diagonal, when i decreases, j must increase. This means within each diagonal, we want elements sorted by increasing j values.

Finally, we need to process diagonals in order. Since diagonal k has sum i + j = k, and we want to process diagonals from smallest to largest sum, we naturally get the diagonal ordering by sorting on i + j.

This leads us to the elegant solution: create tuples (i + j, j, value) for each element. When we sort these tuples:

1. Primary sort by i + j groups elements by diagonal and orders diagonals correctly
2. Secondary sort by j orders elements within each diagonal correctly (smaller j means larger i, which should come first)

The sorting naturally handles both the diagonal grouping and the within-diagonal ordering in one operation.

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

Solution Approach

The implementation follows a sorting-based approach that leverages the diagonal property we identified:

1. Data Collection Phase: We iterate through the entire 2D array using nested loops. For each element at position [i][j] with value v, we create a tuple (i + j, j, v) and append it to an array arr. This transformation encodes both the diagonal information and the ordering within each diagonal.

2. Sorting Phase: We sort the arr array. Python's sort is stable and sorts tuples lexicographically:

   • First by i + j: This groups elements by diagonal since all elements on the same diagonal have the same sum
   • Then by j: Within each diagonal (same i + j), elements are ordered by increasing j values, which gives us the bottom-left to top-right traversal order

3. Extraction Phase: After sorting, we extract just the values (third element of each tuple) using a list comprehension: [v[2] for v in arr]. This gives us the final result in diagonal order.

Time Complexity: O(N log N) where N is the total number of elements in the 2D array. We need to iterate through all elements once (O(N)) and then sort them (O(N log N)).

Space Complexity: O(N) for storing the temporary array of tuples.

The beauty of this solution lies in its simplicity - by encoding the diagonal information directly into sortable tuples, we avoid complex index manipulation and let the sorting algorithm handle the traversal order naturally. The key insight is recognizing that (i + j, j) uniquely identifies both which diagonal an element belongs to and its position within that diagonal.

Example Walkthrough

Let's walk through a small example to illustrate the solution approach:

Consider the array:

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

Step 1: Data Collection Phase

We iterate through each element and create tuples (i+j, j, value):

• Row 0:

  • nums[0][0] = 1: tuple (0+0, 0, 1) = (0, 0, 1)
  • nums[0][1] = 2: tuple (0+1, 1, 2) = (1, 1, 2)
  • nums[0][2] = 3: tuple (0+2, 2, 3) = (2, 2, 3)

• Row 1:

  • nums[1][0] = 4: tuple (1+0, 0, 4) = (1, 0, 4)
  • nums[1][1] = 5: tuple (1+1, 1, 5) = (2, 1, 5)

• Row 2:

  • nums[2][0] = 6: tuple (2+0, 0, 6) = (2, 0, 6)
  • nums[2][1] = 7: tuple (2+1, 1, 7) = (3, 1, 7)
  • nums[2][2] = 8: tuple (2+2, 2, 8) = (4, 2, 8)
  • nums[2][3] = 9: tuple (2+3, 3, 9) = (5, 3, 9)

Our unsorted array of tuples:

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

Step 2: Sorting Phase

Sort the tuples first by i+j (diagonal group), then by j (position within diagonal):

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

Let's see how they're grouped:

• Diagonal 0 (i+j=0): (0,0,1) → element 1
• Diagonal 1 (i+j=1): (1,0,4), (1,1,2) → elements 4, 2 (j=0 comes before j=1)
• Diagonal 2 (i+j=2): (2,0,6), (2,1,5), (2,2,3) → elements 6, 5, 3
• Diagonal 3 (i+j=3): (3,1,7) → element 7
• Diagonal 4 (i+j=4): (4,2,8) → element 8
• Diagonal 5 (i+j=5): (5,3,9) → element 9

Step 3: Extraction Phase

Extract the values (third element) from each sorted tuple:

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

This gives us our final result in diagonal order, with each diagonal traversed from bottom-left to top-right!

Solution Implementation

Time and Space Complexity

The time complexity is O(n × log n), where n is the total number of elements across all rows in the 2D array nums. The algorithm iterates through all elements once to create tuples (taking O(n) time), then sorts the array of tuples (taking O(n × log n) time), and finally extracts the values in a list comprehension (taking O(n) time). The sorting operation dominates, resulting in an overall time complexity of O(n × log n).

The space complexity is O(n), where n is the total number of elements in nums. The algorithm creates an array arr that stores one tuple (i + j, j, v) for each element in the input, requiring O(n) space. The output list also requires O(n) space, but since we typically don't count the output space in space complexity analysis, the overall space complexity remains O(n).

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

Common Pitfalls

Pitfall 1: Incorrect Diagonal Ordering Within Each Diagonal

The Problem: A common mistake is sorting by (i + j, i, value) instead of (i + j, j, value). This would traverse each diagonal from top-right to bottom-left instead of the required bottom-left to top-right order.

Why it happens: It's intuitive to think about sorting by row index i since we often process arrays row by row. However, within a diagonal, as we move from bottom-left to top-right, the row index i decreases while column index j increases.

Example of incorrect approach:

# INCORRECT - sorts by row index within diagonal
diagonal_elements.append((row_index + col_index, row_index, value))

For the example [[1,2,3], [4,5], [6,7,8,9]], this would produce:

• Second diagonal would be [2, 4] instead of [4, 2]
• Third diagonal would be [3, 5, 6] instead of [6, 5, 3]

Solution: Always use column index j as the secondary sort key:

# CORRECT - sorts by column index within diagonal
diagonal_elements.append((row_index + col_index, col_index, value))

Pitfall 2: Assuming Rectangular Matrix

The Problem: Treating the input as a rectangular matrix and using len(nums[0]) to determine the number of columns, which fails for jagged arrays where rows have different lengths.

Example of incorrect approach:

# INCORRECT - assumes all rows have same length
for i in range(len(nums)):
    for j in range(len(nums[0])):  # This assumes all rows are same length!
        if j < len(nums[i]):
            diagonal_elements.append((i + j, j, nums[i][j]))

Solution: Iterate through each row individually using its actual length:

# CORRECT - handles jagged arrays
for row_index, row in enumerate(nums):
    for col_index, value in enumerate(row):
        diagonal_elements.append((row_index + col_index, col_index, value))

Pitfall 3: Memory-Inefficient Alternative Approaches

The Problem: Trying to pre-compute the maximum number of diagonals and using a 2D list to group elements by diagonal, which can waste memory and add complexity.

Example of inefficient approach:

# INEFFICIENT - creates many intermediate lists
max_diagonals = len(nums) + max(len(row) for row in nums) - 1
diagonals = [[] for _ in range(max_diagonals)]
for i in range(len(nums)):
    for j in range(len(nums[i])):
        diagonals[i + j].append(nums[i][j])
# Then need to reverse each diagonal and concatenate...

Solution: The sorting approach with tuples is more elegant and avoids creating intermediate diagonal groupings, using only a single list of tuples that gets sorted in-place.