1329 - Sort the Matrix Diagonally

Difficulty: Medium | Pattern: Matrix + Sorting | Company tags: Amazon, Google

Problem Statement

A matrix diagonal is a diagonal line of cells starting from some cell in either the topmost row or leftmost column and going in the bottom-right direction until reaching the matrix's end.

Given an m x n matrix mat of integers, sort each matrix diagonal in ascending order and return the resulting matrix.

Example:

Input: mat = [[3,3,1,1],[2,2,1,2],[1,1,1,2]]
Output: [[1,1,1,1],[1,2,2,2],[1,2,3,3]]

Approach: Group by Diagonal + Sort — O(m×n × log(min(m,n)))

Key insight: All cells on the same diagonal share the same value of row - col. Group cells by this key, sort each group, and write back.

from collections import defaultdict

def diagonalSort(mat: list[list[int]]) -> list[list[int]]:
    m, n = len(mat), len(mat[0])
    diag = defaultdict(list)
    
    # Collect elements by diagonal
    for r in range(m):
        for c in range(n):
            diag[r - c].append(mat[r][c])
    
    # Sort each diagonal
    for key in diag:
        diag[key].sort()
    
    # Write back
    diag_idx = defaultdict(int)
    for r in range(m):
        for c in range(n):
            key = r - c
            mat[r][c] = diag[key][diag_idx[key]]
            diag_idx[key] += 1
    
    return mat

Algorithm Flow

Dry Run

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

Diagonals (r-c):

• r-c = -2: [mat[0][2], mat[1][3]] = [1,2] → sorted: [1,2]
• r-c = -1: [mat[0][1], mat[1][2], mat[2][3]] = [3,1,2] → sorted: [1,2,3]
• r-c = 0: [mat[0][0], mat[1][1], mat[2][2]] = [3,2,1] → sorted: [1,2,3]
• r-c = 1: [mat[1][0], mat[2][1]] = [2,1] → sorted: [1,2]
• r-c = 2: [mat[2][0]] = [1] → [1]

Write back column by column gives [[1,1,1,1],[1,2,2,2],[1,2,3,3]] ✓

Complexity

• Time: O(m×n × log(min(m,n))) — collecting O(m×n), sorting each diagonal O(d log d)
• Space: O(m×n) for storing diagonal elements

Key Terms

Term	Definition
Diagonal key (row - col)	A constant value shared by all cells on the same top-left-to-bottom-right diagonal; used to group cells without explicit traversal logic.
Hash map grouping	Using a defaultdict(list) to bucket matrix cells by diagonal key before sorting each bucket independently.
In-place write-back	Refilling the original matrix using a per-key read index after sorting, avoiding allocation of a second full matrix.
Row-major iteration	Traversing (r, c) in the natural nested-loop order, which guarantees cells within the same diagonal are revisited in top-to-bottom order during write-back.

FAQ

1. Can this be solved without extra space? Not with the standard approach — O(m×n) space is needed to bucket and sort diagonals. Some in-place variants use a min-heap per diagonal or repeated swaps, but they don't reduce the asymptotic space in the general case.
2. What if the matrix has only one row or one column? Each cell forms its own diagonal (since row - col is unique per cell in a single row/column when the other dimension is 1), so sorting is a no-op and the matrix is returned unchanged.
3. How would this change if we had to sort diagonals in descending order instead? Just sort each diag[key] list with reverse=True (or sort ascending and write back in reverse read order); the grouping logic is unchanged.
4. What if we needed to sort anti-diagonals (top-right to bottom-left) instead? Use row + col as the grouping key instead of row - col, since that sum is constant along anti-diagonals.
5. Is there a way to avoid extra space using min-heaps per diagonal? Yes — push each cell's value onto a heap keyed by diagonal, then pop in order during write-back. This has the same O(m×n log(min(m,n))) time complexity but avoids a separate sort call per diagonal; space is still O(m×n) for the heaps.

Quick Revision

• Goal: sort every top-left-to-bottom-right diagonal of a matrix in ascending order.
• Key insight: all cells on the same diagonal share the same row - col value.
• Step 1: bucket every cell's value into diag[row - col] via a defaultdict(list).
• Step 2: sort each bucket independently.
• Step 3: re-traverse the matrix in row-major order, popping the next smallest value from the matching diagonal bucket each time.
• Time: O(m×n log(min(m,n))) — bucketing is O(m×n), sorting a diagonal of length d costs O(d log d).
• Space: O(m×n) to store all diagonal buckets.
• Anti-diagonal variant: swap the grouping key to row + col.
• Common bug: forgetting to track a separate read index per diagonal key during write-back, which would reuse the same sorted value repeatedly.

Related Problems

• Same "group cells by a derived key" trick: 1424 - Diagonal Traverse II
• Diagonal traversal fundamentals: 498 - Diagonal Traverse
• Matrix transformation via grouping/sorting: 56 - Merge Intervals