Problem Description

Given a 2D matrix with m rows and n columns, you need to return all elements of the matrix in diagonal order.

The diagonal traversal follows a zigzag pattern:

• Start from the top-left corner mat[0][0]
• Move diagonally up-right until you hit a boundary
• Then move diagonally down-left until you hit a boundary
• Continue alternating between these two diagonal directions

For example, if you have a 3x3 matrix:

1 2 3
4 5 6
7 8 9

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

The traversal pattern looks like:

• First diagonal (up-right): 1
• Second diagonal (down-left): 2, 4
• Third diagonal (up-right): 7, 5, 3
• Fourth diagonal (down-left): 6, 8
• Fifth diagonal (up-right): 9

The solution works by iterating through each diagonal line (there are m + n - 1 diagonals total). For each diagonal k:

• It determines the starting position (i, j) based on whether we're still in the first n diagonals or beyond
• It collects all elements along that diagonal from top-right to bottom-left
• If k is even (0, 2, 4...), it reverses the collected elements to achieve the zigzag pattern
• All collected elements are added to the final result array

Intuition

The key observation is that elements on the same diagonal share a common property: the sum of their row and column indices is constant. For example, elements mat[0][2], mat[1][1], and mat[2][0] all have i + j = 2.

This gives us the insight that we can group elements by their diagonal number k, where k = i + j. Since the matrix has m rows and n columns, the diagonal sums range from 0 (top-left corner) to m + n - 2 (bottom-right corner), giving us exactly m + n - 1 diagonals.

For each diagonal k, we need to find all valid (i, j) pairs where i + j = k. The constraints are:

• 0 <= i < m (within row bounds)
• 0 <= j < n (within column bounds)
• i + j = k

From these constraints, we can derive that for diagonal k:

• If k < n, we start from the top row (i = 0) and column j = k
• If k >= n, we've exceeded the rightmost column, so we start from the rightmost column (j = n - 1) and row i = k - n + 1

We then traverse diagonally by incrementing i and decrementing j until we hit a boundary.

The zigzag pattern emerges from alternating the direction of each diagonal. Notice that:

• Even-numbered diagonals (k = 0, 2, 4...) go upward (from bottom-left to top-right)
• Odd-numbered diagonals (k = 1, 3, 5...) go downward (from top-right to bottom-left)

Since we always collect elements from top-right to bottom-left, we need to reverse the elements for even-numbered diagonals to achieve the upward direction.

Solution Approach

The implementation uses a fixed point traversal strategy for each diagonal. Here's how the algorithm works step by step:

1. Initialize variables: Get the dimensions m (rows) and n (columns) of the matrix, and create an empty result list ans.

2. Iterate through all diagonals: Loop through k from 0 to m + n - 2 (inclusive), representing each diagonal.

3. Determine starting position for diagonal k:

   • If k < n: Start from the top row, so i = 0 and j = k
   • If k >= n: We've passed the last column, so start from the rightmost column with i = k - n + 1 and j = n - 1

4. Collect elements along the diagonal: Create a temporary list t to store elements on the current diagonal. Starting from position (i, j), traverse diagonally by:

   • Adding mat[i][j] to the temporary list
   • Moving to the next position: i += 1 (move down) and j -= 1 (move left)
   • Continue while i < m and j >= 0 (stay within matrix bounds)

5. Handle zigzag pattern:

   • If k is even (0, 2, 4...), reverse the temporary list t using t[::-1]
   • This reversal converts the down-left traversal into an up-right direction for the final result

6. Append to result: Extend the answer list with all elements from the current diagonal using ans.extend(t)

7. Return the result: After processing all diagonals, return the complete ans list containing all matrix elements in diagonal order.

The algorithm has a time complexity of O(m * n) since we visit each element exactly once, and a space complexity of O(1) excluding the output array (or O(min(m, n)) if we count the temporary list t which stores at most one diagonal's worth of elements).

Example Walkthrough

Let's walk through a small 3×2 matrix to illustrate the solution approach:

Matrix:
1 2
3 4
5 6

Step 1: Initialize

• m = 3 (rows), n = 2 (columns)
• Total diagonals = m + n - 1 = 4
• Result list ans = []

Step 2: Process diagonal k = 0

• Since k < n (0 < 2), start at: i = 0, j = 0
• Collect elements along diagonal: t = [1]
• Since k is even, reverse: t = [1] (no change for single element)
• Add to result: ans = [1]

Step 3: Process diagonal k = 1

• Since k < n (1 < 2), start at: i = 0, j = 1
• Collect elements from (0,1) to (1,0):
  • Add mat[0][1] = 2 → t = [2]
  • Move to (1,0): Add mat[1][0] = 3 → t = [2, 3]
  • Move to (2,-1): j < 0, stop
• Since k is odd, don't reverse: t = [2, 3]
• Add to result: ans = [1, 2, 3]

Step 4: Process diagonal k = 2

• Since k >= n (2 >= 2), start at: i = k - n + 1 = 1, j = n - 1 = 1
• Collect elements from (1,1) to (2,0):
  • Add mat[1][1] = 4 → t = [4]
  • Move to (2,0): Add mat[2][0] = 5 → t = [4, 5]
  • Move to (3,-1): i >= m, stop
• Since k is even, reverse: t = [5, 4]
• Add to result: ans = [1, 2, 3, 5, 4]

Step 5: Process diagonal k = 3

• Since k >= n (3 >= 2), start at: i = k - n + 1 = 2, j = n - 1 = 1
• Collect elements from (2,1):
  • Add mat[2][1] = 6 → t = [6]
  • Move to (3,0): i >= m, stop
• Since k is odd, don't reverse: t = [6]
• Add to result: ans = [1, 2, 3, 5, 4, 6]

Final Result: [1, 2, 3, 5, 4, 6]

The traversal pattern creates a zigzag:

• Diagonal 0 (even): ↗ moves up-right: [1]
• Diagonal 1 (odd): ↙ moves down-left: [2, 3]
• Diagonal 2 (even): ↗ moves up-right: [5, 4]
• Diagonal 3 (odd): ↙ moves down-left: [6]

Solution Implementation

Time and Space Complexity

The time complexity is O(m × n), where m is the number of rows and n is the number of columns in the matrix. This is because:

• The outer loop runs m + n - 1 times (total number of diagonals)
• For each diagonal k, the inner while loop processes elements along that diagonal
• The total number of elements processed across all diagonals is exactly m × n (every element in the matrix is visited exactly once)
• The reversal operation t[::-1] takes O(len(t)) time, but the sum of all diagonal lengths equals m × n

The space complexity is O(min(m, n)) for the temporary list t that stores elements of each diagonal. The longest diagonal has length min(m, n), which is the maximum space used by t at any point. However, if we ignore the space used for the answer array ans (as mentioned in the reference), and only consider auxiliary space, the space complexity can be considered O(min(m, n)). The reference answer states O(1) which would be accurate if we optimize the code to directly append to the result without using the temporary list t, but with the current implementation using temporary storage t, it's technically O(min(m, n)).

Common Pitfalls

1. Incorrect Starting Position Calculation

One of the most common mistakes is incorrectly determining the starting position for each diagonal, especially when transitioning from diagonals that start on the top row to those that start on the rightmost column.

Pitfall Example:

# Incorrect approach - off by one error
start_row = 0 if diagonal_index <= cols else diagonal_index - cols
# This would cause index out of bounds for later diagonals

Solution: Ensure the correct formula when diagonal_index >= cols:

start_row = 0 if diagonal_index < cols else diagonal_index - cols + 1

The +1 is crucial because when we reach diagonal_index = cols, we need to start from row 1, not row 0.

2. Boundary Check Errors During Traversal

Another frequent mistake is using incorrect boundary conditions while traversing each diagonal, leading to IndexError or missing elements.

Pitfall Example:

# Wrong boundary check
while current_row <= rows and current_col >= 0:  # Should be < rows, not <= rows
    diagonal_elements.append(mat[current_row][current_col])

Solution: Use strict less-than comparison for row bounds:

while current_row < rows and current_col >= 0:
    diagonal_elements.append(mat[current_row][current_col])

3. Incorrect Zigzag Pattern Implementation

Developers often reverse the wrong diagonals or forget to implement the zigzag pattern entirely.

Pitfall Example:

# Reversing odd diagonals instead of even
if diagonal_index % 2 == 1:  # Wrong - should reverse even diagonals
    diagonal_elements = diagonal_elements[::-1]

Solution: Remember that we reverse even-indexed diagonals (0, 2, 4...) to achieve the up-right traversal effect:

if diagonal_index % 2 == 0:
    diagonal_elements = diagonal_elements[::-1]

4. Edge Case: Single Row or Column Matrix

The algorithm might fail or produce incorrect results for matrices with only one row or one column if not handled properly.

Pitfall Example: Not considering that a 1×n or m×1 matrix still needs proper diagonal traversal.

Solution: The provided algorithm naturally handles these cases, but always test with:

• Single row matrix: [[1, 2, 3]] → [1, 2, 3]
• Single column matrix: [[1], [2], [3]] → [1, 2, 3]
• Single element matrix: [[1]] → [1]

5. Inefficient Memory Usage

Creating unnecessary intermediate data structures or not properly managing the temporary diagonal list.

Pitfall Example:

# Creating a new list for every append operation
result = result + diagonal_elements  # Creates new list each time

Solution: Use extend() method for better performance:

result.extend(diagonal_elements)  # Modifies list in place