Problem Description

You are given an m x n matrix of integers and need to sort each diagonal of the matrix in ascending order.

A matrix diagonal is a line of cells that starts from either:

• Any cell in the topmost row (row 0), OR
• Any cell in the leftmost column (column 0)

And proceeds diagonally down-right until it reaches the edge of the matrix.

For example, in a 6 x 3 matrix, the diagonal starting from mat[2][0] would include the cells mat[2][0], mat[3][1], and mat[4][2].

Your task is to:

1. Identify all diagonals in the matrix
2. Sort each diagonal independently in ascending order
3. Return the modified matrix with all diagonals sorted

The key insight is that elements on the same diagonal share a common property: for any two elements (i₁, j₁) and (i₂, j₂) on the same diagonal, the difference j - i is constant. This property helps identify which elements belong to the same diagonal.

The solution uses this property by grouping elements with the same m - i + j value (where m is added as an offset to ensure positive indices), sorting each group, and then placing the sorted elements back into their diagonal positions in the matrix.

Intuition

The first observation is that we need to identify which elements belong to the same diagonal. If we look at any diagonal, we notice that as we move down-right (incrementing row by 1 and column by 1), the difference between column index and row index remains constant.

For example, consider elements on a diagonal:

• mat[0][2]: difference is 2 - 0 = 2
• mat[1][3]: difference is 3 - 1 = 2
• mat[2][4]: difference is 4 - 2 = 2

This pattern holds for every diagonal! Each diagonal can be uniquely identified by its j - i value.

Now, instead of sorting each diagonal in-place (which would be complex to implement), we can:

1. Extract all elements from each diagonal into separate arrays
2. Sort these arrays
3. Put the sorted elements back

To group elements by diagonal, we use the formula m - i + j where m is the number of rows. The addition of m ensures all indices are positive (since j - i can be negative when i > j).

The clever part of the solution is using a list of lists g where g[k] contains all elements from the diagonal with identifier k. After collecting all elements:

• We sort each diagonal's elements
• We iterate through the matrix again and refill each position with the sorted values
• Using pop() ensures we place elements in the correct order as we traverse the matrix

This approach transforms a complex 2D sorting problem into multiple simple 1D sorting problems, making the solution both elegant and efficient.

Learn more about Sorting patterns.

Solution Approach

The implementation follows these steps:

1. Initialize data structures: Create a list g with m + n empty lists. Each list will store elements from one diagonal. The size m + n is sufficient because there are at most m + n - 1 diagonals in an m x n matrix.

2. Group elements by diagonal: Iterate through the matrix and add each element to its corresponding diagonal group:

   for i, row in enumerate(mat):
       for j, x in enumerate(row):
           g[m - i + j].append(x)

   The index m - i + j uniquely identifies each diagonal. Elements with the same m - i + j value belong to the same diagonal.

3. Sort each diagonal: Sort each list in g in reverse order (descending):

   for e in g:
       e.sort(reverse=True)

   We sort in reverse because we'll use pop() later, which removes from the end of the list. This ensures we get elements in ascending order when popping.

4. Refill the matrix: Traverse the matrix again in the same order and replace each element with the sorted values:

   for i in range(m):
       for j in range(n):
           mat[i][j] = g[m - i + j].pop()

   Since we traverse in the same order as we collected elements, and we're popping from lists sorted in reverse, each diagonal gets filled with its elements in ascending order.

The time complexity is O(m * n * log(min(m, n))) because each element is part of a diagonal of maximum length min(m, n), and sorting each diagonal takes O(k * log(k)) where k is the diagonal length. The space complexity is O(m * n) for storing all elements in the auxiliary lists.

Example Walkthrough

Let's walk through a small example with a 3 x 4 matrix:

Initial matrix:
[3, 3, 1, 1]
[2, 2, 1, 2]
[1, 1, 1, 2]

Step 1: Identify the diagonals

Each diagonal is identified by m - i + j (where m = 3):

• Position [0][0]: 3 - 0 + 0 = 3, value = 3
• Position [0][1]: 3 - 0 + 1 = 4, value = 3
• Position [0][2]: 3 - 0 + 2 = 5, value = 1
• Position [0][3]: 3 - 0 + 3 = 6, value = 1
• Position [1][0]: 3 - 1 + 0 = 2, value = 2
• Position [1][1]: 3 - 1 + 1 = 3, value = 2
• Position [1][2]: 3 - 1 + 2 = 4, value = 1
• Position [1][3]: 3 - 1 + 3 = 5, value = 2
• Position [2][0]: 3 - 2 + 0 = 1, value = 1
• Position [2][1]: 3 - 2 + 1 = 2, value = 1
• Position [2][2]: 3 - 2 + 2 = 3, value = 1
• Position [2][3]: 3 - 2 + 3 = 4, value = 2

Step 2: Group elements by diagonal

After grouping, our g array looks like:

• g[1]: [1] (diagonal starting at [2][0])
• g[2]: [2, 1] (diagonal starting at [1][0])
• g[3]: [3, 2, 1] (diagonal starting at [0][0])
• g[4]: [3, 1, 2] (diagonal starting at [0][1])
• g[5]: [1, 2] (diagonal starting at [0][2])
• g[6]: [1] (diagonal starting at [0][3])

Step 3: Sort each diagonal in reverse

After sorting in reverse order:

• g[1]: [1]
• g[2]: [2, 1]
• g[3]: [3, 2, 1]
• g[4]: [3, 2, 1]
• g[5]: [2, 1]
• g[6]: [1]

Step 4: Refill the matrix

Traverse the matrix in the same order and pop from each diagonal:

• [0][0]: pop from g[3] → 1 (g[3] becomes [3, 2])
• [0][1]: pop from g[4] → 1 (g[4] becomes [3, 2])
• [0][2]: pop from g[5] → 1 (g[5] becomes [2])
• [0][3]: pop from g[6] → 1 (g[6] becomes [])
• [1][0]: pop from g[2] → 1 (g[2] becomes [2])
• [1][1]: pop from g[3] → 2 (g[3] becomes [3])
• [1][2]: pop from g[4] → 2 (g[4] becomes [3])
• [1][3]: pop from g[5] → 2 (g[5] becomes [])
• [2][0]: pop from g[1] → 1 (g[1] becomes [])
• [2][1]: pop from g[2] → 2 (g[2] becomes [])
• [2][2]: pop from g[3] → 3 (g[3] becomes [])
• [2][3]: pop from g[4] → 3 (g[4] becomes [])

Final sorted matrix:

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

Each diagonal is now sorted in ascending order!

Solution Implementation

Time and Space Complexity

Time Complexity: O(m × n × log(min(m, n)))

The algorithm works by grouping elements along each diagonal, sorting them, and then placing them back. Breaking down the time complexity:

1. First nested loop: Iterating through all m × n elements to group them by diagonal takes O(m × n) time.

2. Sorting phase: There are m + n - 1 diagonals in total. Each diagonal contains at most min(m, n) elements. Sorting each diagonal takes O(k × log k) where k is the diagonal length. The total sorting time across all diagonals is dominated by O(m × n × log(min(m, n))) since the sum of all diagonal lengths equals m × n.

3. Second nested loop: Redistributing the sorted elements back to the matrix takes O(m × n) time.

The overall time complexity is O(m × n) + O(m × n × log(min(m, n))) + O(m × n) = O(m × n × log(min(m, n))).

Space Complexity: O(m × n)

The algorithm uses a list g containing m + n sublists that collectively store all m × n elements from the matrix. This requires O(m × n) additional space to store all the matrix elements in the diagonal groups.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Diagonal Indexing Formula

A critical pitfall is using the wrong formula to identify diagonals. Many developers intuitively try i + j or i - j as the diagonal identifier, but these don't correctly group all diagonals.

Why it fails:

• Using i - j alone produces negative indices for diagonals below the main diagonal
• Using i + j groups anti-diagonals (top-right to bottom-left) instead of the required top-left to bottom-right diagonals

Correct approach: The formula m - i + j (or alternatively i - j + offset) ensures:

• All indices are positive
• Elements on the same diagonal share the same identifier
• The offset prevents array index out of bounds errors

2. Forgetting to Handle the Pop Order

When using pop() to retrieve elements, forgetting to sort in reverse order is a common mistake:

# WRONG: Sorting in ascending order
for diagonal in diagonal_buckets:
    diagonal.sort()  # This will cause elements to be placed in descending order!

# CORRECT: Sorting in descending order for pop()
for diagonal in diagonal_buckets:
    diagonal.sort(reverse=True)  # Elements will be popped in ascending order

Alternative solution without reverse sorting:

# Use deque for efficient popleft()
from collections import deque

# After sorting normally:
for diagonal in diagonal_buckets:
    diagonal.sort()

# Convert to deque
diagonal_buckets = [deque(d) for d in diagonal_buckets]

# Pop from the left instead
for row_idx in range(rows):
    for col_idx in range(cols):
        diagonal_id = rows - row_idx + col_idx
        mat[row_idx][col_idx] = diagonal_buckets[diagonal_id].popleft()

3. Incorrect Bucket Array Size

Allocating too few buckets causes index out of bounds errors:

# WRONG: Not enough buckets
diagonal_buckets = [[] for _ in range(rows)]  # Missing diagonals!

# CORRECT: Maximum possible diagonals is rows + cols - 1
diagonal_buckets = [[] for _ in range(rows + cols)]  # Safe allocation

The actual number of diagonals is rows + cols - 1, but allocating rows + cols provides a safety margin without significant overhead.