LeetCode 498: Diagonal Traverse

Problem Restatement

We are given an m x n matrix mat.

We need to return all elements of the matrix in diagonal order.

The traversal starts at the top-left cell, then moves along diagonals in alternating directions.

For this matrix:

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

The diagonal order is:

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

The official problem asks to return all elements of an m x n matrix in diagonal order.

Input and Output

Function shape:

def findDiagonalOrder(mat: list[list[int]]) -> list[int]:
    ...

Examples

Example 1:

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

Diagonals by order:

Answer:

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

Example 2:

mat = [
    [1, 2],
    [3, 4],
]

Diagonal order:

[1, 2, 3, 4]

First Thought: Simulate Movement

We can try to move directly through the matrix.

The movement alternates between:

When we hit a boundary, we change direction and move to the next valid starting cell.

This works, but the boundary logic can become easy to get wrong, especially for single-row and single-column matrices.

A simpler method is to process one diagonal at a time.

Key Insight

All cells on the same diagonal have the same value of:

row + col

For example, in a 3 x 3 matrix:

So the matrix has:

m + n - 1

diagonals.

We can collect each diagonal from top-right to bottom-left. Then we reverse every even-numbered diagonal to match the required zigzag order.

Algorithm

Let:

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

For each diagonal index d from 0 to rows + cols - 2:

1. Find the starting cell.
2. Walk down-left while inside the matrix.
3. Store elements of this diagonal in a temporary list.
4. If d is even, reverse the temporary list.
5. Append it to the answer.

Starting cell:

Then move down-left:

row += 1
col -= 1

Correctness

Every cell (row, col) belongs to exactly one diagonal, identified by row + col.

The algorithm processes diagonal indices from 0 to rows + cols - 2, so it covers every possible value of row + col.

For each diagonal d, the chosen starting cell is the topmost or rightmost valid cell on that diagonal. Moving down-left visits exactly all cells whose row and column remain inside the matrix and whose sum is d.

Therefore, every matrix cell is visited exactly once.

The required traversal alternates direction by diagonal. The algorithm collects every diagonal in top-right to bottom-left order. For odd diagonals, this is already the required order. For even diagonals, reversing gives bottom-left to top-right order.

Thus the final list contains all matrix elements in the required diagonal order.

Complexity

The returned answer itself has size O(mn).

Implementation

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

        ans = []

        for d in range(rows + cols - 1):
            if d < cols:
                row = 0
                col = d
            else:
                row = d - cols + 1
                col = cols - 1

            diagonal = []

            while row < rows and col >= 0:
                diagonal.append(mat[row][col])
                row += 1
                col -= 1

            if d % 2 == 0:
                diagonal.reverse()

            ans.extend(diagonal)

        return ans

Code Explanation

The number of diagonals is:

rows + cols - 1

For each diagonal, we compute the starting point:

if d < cols:
    row = 0
    col = d
else:
    row = d - cols + 1
    col = cols - 1

Then we collect cells down-left:

while row < rows and col >= 0:
    diagonal.append(mat[row][col])
    row += 1
    col -= 1

For even diagonals, the required direction is the opposite of how we collected it:

if d % 2 == 0:
    diagonal.reverse()

Then we append this diagonal to the final answer:

ans.extend(diagonal)

Testing

def run_tests():
    s = Solution()

    assert s.findDiagonalOrder([
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9],
    ]) == [1, 2, 4, 7, 5, 3, 6, 8, 9]

    assert s.findDiagonalOrder([
        [1, 2],
        [3, 4],
    ]) == [1, 2, 3, 4]

    assert s.findDiagonalOrder([[1]]) == [1]

    assert s.findDiagonalOrder([[1, 2, 3]]) == [1, 2, 3]

    assert s.findDiagonalOrder([
        [1],
        [2],
        [3],
    ]) == [1, 2, 3]

    print("all tests passed")

run_tests()

Test meaning: