LeetCode 3446 - Sort Matrix by Diagonals

We are given an n x n square matrix grid. Every cell belongs to exactly one top-left to bottom-right diagonal. Two cells (r1, c1) and (r2, c2) are on the same diagonal if r1 - c1 == r2 - c2.

leetcodemediumarraysortingmatrix

LeetCode Problem 3446

Difficulty: 🟡 Medium
Topics: Array, Sorting, Matrix

Solution

Problem Understanding

We are given an n x n square matrix grid. Every cell belongs to exactly one top-left to bottom-right diagonal. Two cells (r1, c1) and (r2, c2) are on the same diagonal if r1 - c1 == r2 - c2.

The problem asks us to independently sort each diagonal, but the required sorting order depends on which side of the matrix the diagonal belongs to.

For diagonals in the bottom-left triangle, including the main diagonal, the values must appear in non-increasing order when traversed from the diagonal's starting position toward the bottom-right. These are precisely the diagonals whose starting point lies in the first column, including (0,0).

For diagonals in the top-right triangle, the values must appear in non-decreasing order. These are the diagonals whose starting point lies in the top row, excluding (0,0).

The input is a square matrix of size n × n, where 1 <= n <= 10. The matrix elements can be negative or positive integers.

The small constraint is important. Even though n is only at most 10, we should still design a clean and efficient solution. Since the matrix contains at most 100 elements, sorting each diagonal independently is very manageable.

Several edge cases are worth noting:

  • A 1 x 1 matrix contains only one diagonal of length one, so no changes are needed.
  • Diagonals of length one are already sorted regardless of the required order.
  • Negative values and duplicate values must be handled correctly by the sorting operation.
  • The main diagonal belongs to the bottom-left group and therefore must be sorted in non-increasing order.

Approaches

Brute Force Approach

A straightforward approach is to repeatedly inspect every diagonal and perform swaps until the diagonal becomes sorted in the required order.

For each diagonal, we could run a bubble sort or another comparison-based in-place sorting process directly on the matrix cells. Since a diagonal may contain up to n elements, and there are 2n - 1 diagonals, repeatedly swapping values would eventually produce the correct ordering.

This approach is correct because each diagonal is processed independently and eventually reaches the desired sorted state. However, using repeated comparisons and swaps is unnecessarily expensive compared to extracting the diagonal into a temporary array, sorting it once, and writing it back.

Key Insight

Every diagonal is completely independent from all others.

A top-left to bottom-right diagonal is uniquely identified by its starting position. Therefore, we can:

  1. Collect all values from one diagonal.
  2. Sort the collected values in the required order.
  3. Write the sorted values back into the same diagonal positions.

Since sorting is already highly optimized, this is simpler and more efficient than manually performing repeated swaps.

The only challenge is determining which order to use:

  • Diagonals starting in the first column belong to the bottom-left region and must be sorted descending.
  • Diagonals starting in the top row, excluding (0,0), belong to the top-right region and must be sorted ascending.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(n³) O(1) Repeatedly swap elements along diagonals until sorted
Optimal O(n² log n) O(n) Extract each diagonal, sort once, write back

Algorithm Walkthrough

Optimal Algorithm

  1. Process every diagonal that starts in the first column.

These diagonals correspond to the bottom-left triangle, including the main diagonal. For each starting row r, begin at (r, 0) and collect all values while moving down-right. 2. Sort the collected values in descending order.

The problem requires all bottom-left diagonals to be non-increasing, so sorting in reverse order directly produces the desired arrangement. 3. Write the sorted values back into the same diagonal positions.

Traverse the diagonal again and replace each cell with the corresponding sorted value. 4. Process every diagonal that starts in the top row except (0,0).

These diagonals belong to the top-right triangle. For each starting column c > 0, collect all values while moving down-right. 5. Sort the collected values in ascending order.

The problem requires these diagonals to be non-decreasing. 6. Write the sorted values back into the matrix.

Traverse the same diagonal again and place the sorted values into their original positions. 7. After all diagonals have been processed, return the modified matrix.

Why it works

Every matrix cell belongs to exactly one top-left to bottom-right diagonal. The algorithm visits every diagonal exactly once and sorts it according to the rule that applies to its region. Since diagonals do not overlap except through their own elements, sorting one diagonal cannot affect the correctness of another. Therefore, after all diagonals have been processed, every diagonal satisfies its required ordering, which is exactly the desired final matrix.

Python Solution

from typing import List

class Solution:
    def sortMatrix(self, grid: List[List[int]]) -> List[List[int]]:
        n = len(grid)

        def process(start_row: int, start_col: int, reverse: bool) -> None:
            diagonal = []

            r, c = start_row, start_col
            while r < n and c < n:
                diagonal.append(grid[r][c])
                r += 1
                c += 1

            diagonal.sort(reverse=reverse)

            r, c = start_row, start_col
            idx = 0
            while r < n and c < n:
                grid[r][c] = diagonal[idx]
                idx += 1
                r += 1
                c += 1

        # Bottom-left triangle + main diagonal
        for row in range(n):
            process(row, 0, True)

        # Top-right triangle
        for col in range(1, n):
            process(0, col, False)

        return grid

The implementation defines a helper function process that handles a single diagonal. The helper first collects all diagonal values into a temporary list, sorts them according to the requested direction, and then writes the values back into the matrix.

The first loop processes all diagonals that begin in the first column. These correspond to the bottom-left region and therefore use descending order.

The second loop processes all diagonals that begin in the top row, excluding the main diagonal. These correspond to the top-right region and therefore use ascending order.

Because every diagonal is visited exactly once, the resulting matrix satisfies all requirements.

Go Solution

import "sort"

func sortMatrix(grid [][]int) [][]int {
	n := len(grid)

	process := func(startRow, startCol int, descending bool) {
		diagonal := []int{}

		r, c := startRow, startCol
		for r < n && c < n {
			diagonal = append(diagonal, grid[r][c])
			r++
			c++
		}

		sort.Ints(diagonal)

		if descending {
			for i, j := 0, len(diagonal)-1; i < j; i, j = i+1, j-1 {
				diagonal[i], diagonal[j] = diagonal[j], diagonal[i]
			}
		}

		r, c = startRow, startCol
		idx := 0
		for r < n && c < n {
			grid[r][c] = diagonal[idx]
			idx++
			r++
			c++
		}
	}

	for row := 0; row < n; row++ {
		process(row, 0, true)
	}

	for col := 1; col < n; col++ {
		process(0, col, false)
	}

	return grid
}

The Go implementation follows the same strategy as the Python solution. Since Go's standard library provides sort.Ints, each diagonal is first sorted in ascending order. For descending diagonals, the sorted slice is reversed in place before being written back into the matrix.

No special handling for integer overflow is required because the algorithm only compares and moves existing values. Go slices are used as temporary storage for each diagonal.

Worked Examples

Example 1

Input:

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

Main diagonal

Positions:

Cell Value
(0,0) 1
(1,1) 8
(2,2) 6

Collected:

[1, 8, 6]

Descending sort:

[8, 6, 1]

Matrix becomes:

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

Diagonal starting at (1,0)

Collected:

[9, 5]

Descending sort:

[9, 5]

No change.

Diagonal starting at (2,0)

Collected:

[4]

No change.

Diagonal starting at (0,1)

Collected:

[7, 2]

Ascending sort:

[2, 7]

Matrix becomes:

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

Diagonal starting at (0,2)

Collected:

[3]

No change.

Final result:

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

Example 2

Input:

[
  [0,1],
  [1,2]
]

Main diagonal:

[0, 2]

Descending sort:

[2, 0]

Matrix:

[
  [2,1],
  [1,0]
]

Remaining diagonals contain one element each.

Final result:

[
  [2,1],
  [1,0]
]

Example 3

Input:

[
  [1]
]

Only one diagonal exists:

[1]

Already sorted.

Output:

[
  [1]
]

Complexity Analysis

Measure Complexity Explanation
Time O(n² log n) Sorting each diagonal, total work bounded by all diagonal lengths
Space O(n) Temporary storage for one diagonal at a time

The matrix contains n² elements. Each diagonal of length k is extracted and sorted in O(k log k) time. Summing over all diagonals gives O(n² log n) in the worst case. Since only one diagonal array is stored at any moment, the auxiliary space is O(n).

Test Cases

from typing import List

s = Solution()

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

assert s.sortMatrix([[0,1],[1,2]]) == [
    [2,1],
    [1,0]
]  # Example 2

assert s.sortMatrix([[1]]) == [
    [1]
]  # Single element matrix

assert s.sortMatrix([[5,4],[3,2]]) == [
    [5,4],
    [3,2]
]  # Already satisfies ordering

assert s.sortMatrix([[-1,-2],[-3,-4]]) == [
    [-1,-2],
    [-3,-4]
]  # Negative values

assert s.sortMatrix([
    [3,3,3],
    [3,3,3],
    [3,3,3]
]) == [
    [3,3,3],
    [3,3,3],
    [3,3,3]
]  # All duplicates

assert s.sortMatrix([
    [9,8,7,6],
    [5,4,3,2],
    [1,0,-1,-2],
    [-3,-4,-5,-6]
]) == [
    [9,2,7,6],
    [5,4,3,8],
    [1,0,-1,-2],
    [-3,-4,-5,-6]
]  # Larger matrix with mixed diagonals
Test Why
Example 1 Verifies both ascending and descending diagonal processing
Example 2 Smallest non-trivial matrix
Example 3 Single-cell matrix
Already ordered matrix Ensures no unnecessary modifications
Negative values Confirms sorting works with negative numbers
All duplicates Verifies handling of equal values
Larger 4×4 matrix Exercises multiple diagonals of varying lengths

Edge Cases

Single Element Matrix

When n = 1, there is only one diagonal consisting of a single element. A careless implementation might attempt unnecessary processing or run into indexing issues. The solution handles this naturally because collecting and sorting a one-element diagonal leaves it unchanged.

Diagonals of Length One

Many diagonals near the matrix borders contain only one element. These diagonals are already sorted regardless of whether ascending or descending order is required. The algorithm still processes them uniformly, collecting one value, sorting it, and writing it back without special-case logic.

Duplicate Values

If a diagonal contains repeated numbers, a buggy implementation might accidentally assume strict ordering instead of non-increasing or non-decreasing ordering. Using standard sorting guarantees that equal values remain grouped correctly and satisfy the required ordering constraints.

Negative Numbers

The matrix values may be negative. Any solution relying on assumptions about positivity could fail. Since the algorithm uses ordinary comparison-based sorting, negative and positive values are handled identically and produce the correct ordering.

Main Diagonal Classification

A common mistake is deciding whether the main diagonal belongs to the bottom-left region or the top-right region. The problem explicitly states that the middle diagonal is included in the bottom-left triangle. By processing all diagonals starting in the first column, including (0,0), the implementation correctly sorts the main diagonal in non-increasing order.