LeetCode 498 - Diagonal Traverse

This problem asks us to traverse a matrix in a specific diagonal pattern and return all elements in the order they are visited. We are given an m x n matrix mat, where m represents the number of rows and n represents the number of columns.

leetcodemediumarraymatrixsimulation

LeetCode Problem 498

Difficulty: 🟡 Medium
Topics: Array, Matrix, Simulation

Solution

Problem Understanding

This problem asks us to traverse a matrix in a specific diagonal pattern and return all elements in the order they are visited.

We are given an m x n matrix mat, where m represents the number of rows and n represents the number of columns. The goal is to return a one dimensional array containing every matrix element, visited in a diagonal zigzag order.

The traversal follows these rules:

  • Elements are visited diagonal by diagonal.

  • Each diagonal contains cells where the sum of row and column indices is constant.

  • The direction alternates between diagonals:

  • One diagonal is traversed upward-right.

  • The next diagonal is traversed downward-left.

For example, in:

1 2 3
4 5 6
7 8 9

The traversal becomes:

[1]          -> upward
[2,4]        -> downward
[7,5,3]      -> upward
[6,8]        -> downward
[9]          -> upward

Result:

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

The constraints are important:

  • 1 <= m, n <= 10^4
  • 1 <= m * n <= 10^4

The total number of elements is at most 10^4, which suggests that an O(m × n) solution is ideal and fully acceptable. Since every element must appear in the output, we cannot do better than linear time in terms of total cells visited.

A few important edge cases stand out immediately:

A matrix with only one row must simply be traversed left to right. A matrix with only one column must be traversed top to bottom. Rectangular matrices, where m != n, require careful boundary handling because diagonals terminate at different edges. Small matrices like 1 x 1 are also easy places for off by one mistakes.

The problem guarantees the matrix is non-empty and rectangular, meaning every row has the same number of columns.

Approaches

Brute Force Approach

A straightforward way to solve this problem is to explicitly group matrix elements by diagonal.

The key observation is that every cell (row, col) belongs to a diagonal identified by:

row + col

We could iterate through the matrix and store elements in a hash map where the key is row + col. After collecting all diagonals, we iterate through them in order.

For even numbered diagonals, we reverse the collected elements because traversal should move upward-right. For odd numbered diagonals, we keep the order as is because traversal should move downward-left.

This works because all cells sharing the same row + col belong to the same diagonal.

However, this approach requires additional storage for the diagonal groups and extra work for reversing lists. While still efficient enough for constraints, it uses unnecessary auxiliary space.

Key Insight for an Optimal Solution

Instead of storing diagonals separately, we can directly simulate the traversal.

At any moment, we know:

  • Our current position (row, col)
  • The current movement direction

There are only two directions:

  1. Upward-right: (-1, +1)
  2. Downward-left: (+1, -1)

Whenever we hit a boundary, we switch direction and reposition carefully:

  • If moving upward-right:

  • Hitting the right wall means move down.

  • Hitting the top wall means move right.

  • If moving downward-left:

  • Hitting the bottom wall means move right.

  • Hitting the left wall means move down.

This avoids storing diagonals entirely and processes each element exactly once.

Approach Time Complexity Space Complexity Notes
Brute Force O(m × n) O(m × n) Group elements by diagonal and reverse alternating groups
Optimal O(m × n) O(1) excluding output Simulate movement directly with boundary handling

Algorithm Walkthrough

  1. Initialize matrix dimensions.

Store the number of rows and columns:

rows = len(mat)
cols = len(mat[0])

These values are needed for boundary checks during traversal. 2. Prepare result storage.

Create an empty result array that will hold elements in traversal order. 3. Start traversal at the top-left corner.

Initialize:

row = 0
col = 0

This is always the first element. 4. Track movement direction.

Use a variable to represent direction:

direction = 1

Where:

  • 1 means upward-right
  • -1 means downward-left
  1. Process exactly m × n elements.

Since every cell must be visited once, loop exactly rows * cols times.

During each iteration:

  • Append the current element to the result.
  • Compute tentative next position.
  1. Handle upward-right movement.

If moving upward-right:

row -= 1
col += 1

If we go out of bounds:

  • Right boundary:
col == cols

Move down one row and reverse direction.

  • Top boundary:
row < 0

Reset to first valid row and reverse direction. 7. Handle downward-left movement.

If moving downward-left:

row += 1
col -= 1

If we go out of bounds:

  • Bottom boundary:
row == rows

Move right and reverse direction.

  • Left boundary:
col < 0

Reset to first valid column and reverse direction. 8. Continue until all elements are visited.

Since each iteration adds exactly one element, after m × n iterations the result is complete.

Why it works

The algorithm maintains the invariant that (row, col) always points to the next valid matrix cell. Every movement follows the required diagonal direction, and whenever a boundary is encountered, the algorithm transitions to the correct starting point for the next diagonal while reversing direction. Because every cell is visited exactly once and traversal order matches the diagonal zigzag specification, the final output is correct.

Python Solution

from typing import List

class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        rows = len(mat)
        cols = len(mat[0])

        result = []
        row = 0
        col = 0
        direction = 1  # 1 = upward-right, -1 = downward-left

        for _ in range(rows * cols):
            result.append(mat[row][col])

            if direction == 1:
                row -= 1
                col += 1

                if col == cols:
                    col = cols - 1
                    row += 2
                    direction = -1
                elif row < 0:
                    row = 0
                    direction = -1

            else:
                row += 1
                col -= 1

                if row == rows:
                    row = rows - 1
                    col += 2
                    direction = 1
                elif col < 0:
                    col = 0
                    direction = 1

        return result

The implementation begins by determining matrix dimensions and initializing the traversal state. The result list stores elements in visitation order, while row and col track the current position.

The direction variable controls movement. A value of 1 represents upward-right motion, while -1 represents downward-left motion.

The loop executes exactly rows * cols times, guaranteeing that every matrix element is processed once. At each iteration, the current cell is appended to result.

After adding an element, the algorithm updates coordinates according to the current direction. If movement exits matrix bounds, boundary correction logic moves the traversal to the correct next starting point and flips direction.

For example, when moving upward-right and the traversal exits through the right wall, we reposition by stepping down and switching to downward-left movement. Similar corrections exist for all four boundaries.

This direct simulation avoids storing diagonals explicitly and achieves optimal efficiency.

Go Solution

func findDiagonalOrder(mat [][]int) []int {
	rows := len(mat)
	cols := len(mat[0])

	result := make([]int, 0, rows*cols)

	row, col := 0, 0
	direction := 1 // 1 = upward-right, -1 = downward-left

	for i := 0; i < rows*cols; i++ {
		result = append(result, mat[row][col])

		if direction == 1 {
			row--
			col++

			if col == cols {
				col = cols - 1
				row += 2
				direction = -1
			} else if row < 0 {
				row = 0
				direction = -1
			}
		} else {
			row++
			col--

			if row == rows {
				row = rows - 1
				col += 2
				direction = 1
			} else if col < 0 {
				col = 0
				direction = 1
			}
		}
	}

	return result
}

The Go implementation follows exactly the same traversal logic as the Python version.

A slice is preallocated using:

make([]int, 0, rows*cols)

This avoids repeated reallocations as elements are appended.

Unlike Python lists, Go slices benefit from explicit capacity reservation for efficiency. Integer overflow is not a concern here because the maximum number of matrix elements is only 10^4.

Since LeetCode guarantees a valid non-empty matrix, no special handling for empty input is required.

Worked Examples

Example 1

Input:

mat = [
  [1,2,3],
  [4,5,6],
  [7,8,9]
]
Step Position Value Direction Result
1 (0,0) 1 Up [1]
2 (0,1) 2 Down [1,2]
3 (1,0) 4 Down [1,2,4]
4 (2,0) 7 Up [1,2,4,7]
5 (1,1) 5 Up [1,2,4,7,5]
6 (0,2) 3 Up [1,2,4,7,5,3]
7 (1,2) 6 Down [1,2,4,7,5,3,6]
8 (2,1) 8 Down [1,2,4,7,5,3,6,8]
9 (2,2) 9 Up [1,2,4,7,5,3,6,8,9]

Final output:

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

Example 2

Input:

mat = [
  [1,2],
  [3,4]
]
Step Position Value Direction Result
1 (0,0) 1 Up [1]
2 (0,1) 2 Down [1,2]
3 (1,0) 3 Down [1,2,3]
4 (1,1) 4 Up [1,2,3,4]

Final output:

[1,2,3,4]

Complexity Analysis

Measure Complexity Explanation
Time O(m × n) Every matrix cell is visited exactly once
Space O(1) excluding output Only a few variables are used regardless of matrix size

The algorithm processes each element exactly once, making the runtime proportional to the total number of matrix entries.

The auxiliary space complexity is constant because only a fixed number of variables are maintained for traversal state. The output array itself requires O(m × n) space, but this is typically excluded from auxiliary space analysis.

Test Cases

solution = Solution()

# Example 1
assert solution.findDiagonalOrder(
    [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
) == [1, 2, 4, 7, 5, 3, 6, 8, 9]  # standard square matrix

# Example 2
assert solution.findDiagonalOrder(
    [[1, 2], [3, 4]]
) == [1, 2, 3, 4]  # small even-sized matrix

# Single element
assert solution.findDiagonalOrder(
    [[42]]
) == [42]  # minimum valid input

# Single row
assert solution.findDiagonalOrder(
    [[1, 2, 3, 4]]
) == [1, 2, 3, 4]  # horizontal traversal only

# Single column
assert solution.findDiagonalOrder(
    [[1], [2], [3], [4]]
) == [1, 2, 3, 4]  # vertical traversal only

# Rectangular matrix, more rows
assert solution.findDiagonalOrder(
    [[1, 2], [3, 4], [5, 6]]
) == [1, 2, 3, 5, 4, 6]  # non-square matrix

# Rectangular matrix, more columns
assert solution.findDiagonalOrder(
    [[1, 2, 3, 4], [5, 6, 7, 8]]
) == [1, 2, 5, 6, 3, 4, 7, 8]  # wide matrix

# Negative values
assert solution.findDiagonalOrder(
    [[-1, -2], [-3, -4]]
) == [-1, -2, -3, -4]  # handles negative numbers

# Larger matrix
assert solution.findDiagonalOrder(
    [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9],
        [10, 11, 12]
    ]
) == [1, 2, 4, 7, 5, 3, 6, 8, 10, 11, 9, 12]  # multi-diagonal case
Test Why
[[1,2,3],[4,5,6],[7,8,9]] Verifies standard square traversal
[[1,2],[3,4]] Validates small matrix behavior
[[42]] Tests minimum input size
[[1,2,3,4]] Ensures single row handling
[[1],[2],[3],[4]] Ensures single column handling
[[1,2],[3,4],[5,6]] Verifies tall rectangular matrices
[[1,2,3,4],[5,6,7,8]] Verifies wide rectangular matrices
Negative values Confirms values themselves do not affect traversal
Larger matrix Tests multiple direction changes

Edge Cases

Single Row Matrix

A matrix with only one row can easily break diagonal traversal logic because upward movement immediately hits boundaries. For example:

[[1,2,3,4]]

The expected output is simply left to right traversal. The implementation handles this naturally because each upward-right move immediately triggers top boundary correction.

Single Column Matrix

A single column matrix introduces the opposite challenge:

[[1],
 [2],
 [3]]

Downward-left movement quickly exits through the left boundary. Incorrect boundary handling can cause index errors or repeated elements. The implementation correctly resets the column to 0 and switches direction.

Rectangular Matrices

Non-square matrices are a common source of bugs because many implementations implicitly assume equal dimensions. For example:

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

Some diagonals terminate at the right edge while others terminate at the bottom edge. The implementation explicitly checks all four matrix boundaries independently, ensuring correct traversal regardless of shape.

One Element Matrix

The smallest valid input:

[[5]]

can expose off by one mistakes. Since the loop executes exactly m × n times, the implementation processes the single value once and terminates cleanly without invalid movement.