Problem

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. For example, the matrix diagonal starting from mat[2][0], where mat is a 6 x 3 matrix, includes cells mat[2][0], mat[3][1], and mat[4][2].

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

Examples

Example 1:

$$ Input: \begin{array}{cccc} \colorbox{yellow}3 & \colorbox{lightgray}3 & \colorbox{cyan}1 & \colorbox{orange}1 \\ \colorbox{lightgreen}2 & \colorbox{yellow}2 & \colorbox{lightgray}1 & \colorbox{cyan}2 \\ \colorbox{lightblue}1 & \colorbox{lightgreen}1 & \colorbox{yellow}1 & \colorbox{lightgray}2 \\ \end{array} \\ , Output: \begin{array}{cccc} \colorbox{yellow}1 & \colorbox{lightgray}1 & \colorbox{cyan}1 & \colorbox{orange}1 \\ \colorbox{lightgreen}1 & \colorbox{yellow}2 & \colorbox{lightgray}2 & \colorbox{cyan}2 \\ \colorbox{lightblue}1 & \colorbox{lightgreen}2 & \colorbox{yellow}3 & \colorbox{lightgray}3 \\ \end{array} $$

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]]

Example 2:

Input: mat = [[11,25,66,1,69,7],[23,55,17,45,15,52],[75,31,36,44,58,8],[22,27,33,25,68,4],[84,28,14,11,5,50]]
Output: [[5,17,4,1,52,7],[11,11,25,45,8,69],[14,23,25,44,58,15],[22,27,31,36,50,66],[84,28,75,33,55,68]]

Solution

Method 1 - Iteration

To solve this problem, you need to sort each diagonal of the matrix individually. A diagonal in the matrix starts from the topmost row or the leftmost column and extends diagonally to the bottom-right until the matrix’s end.

Approach

  1. Traverse each diagonal starting from the top row.
  2. Traverse each diagonal starting from the left column (excluding the top-left corner since it’s already covered in the first step).
  3. Extract each diagonal into a list, sort it, and place it back into the matrix.

Code

Java
public class Solution {
    public int[][] diagonalSort(int[][] mat) {
        int m = mat.length, n = mat[0].length;

        // Traverse the diagonals originating from the first row
        for (int col = 0; col < n; col++) {
            sortDiagonal(mat, 0, col, m, n);
        }
        
        // Traverse the diagonals originating from the first column
        for (int row = 1; row < m; row++) {
            sortDiagonal(mat, row, 0, m, n);
        }
        
        return mat;
    }
    
    private void sortDiagonal(int[][] mat, int row, int col, int m, int n) {
        List<Integer> diag = new ArrayList<>();
        int r = row, c = col;
        
        // Extract the diagonal elements
        while (r < m && c < n) {
            diag.add(mat[r][c]);
            r++;
            c++;
        }
        
        // Sort the diagonal elements
        Collections.sort(diag);
        
        // Place them back in the matrix
        r = row;
        c = col;
        int index = 0;
        while (r < m && c < n) {
            mat[r][c] = diag.get(index++);
            r++;
            c++;
        }
    }
}
Python
class Solution:
    def diagonalSort(self, mat: List[List[int]]) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        diagonals = defaultdict(list)
        
        # Gather all elements of each diagonal
        for i in range(m):
            for j in range(n):
                diagonals[i - j].append(mat[i][j])
                
        # Sort each diagonal
        for key in diagonals:
            diagonals[key].sort()
            
        # Put sorted elements back into their respective diagonals
        for i in range(m):
            for j in range(n):
                mat[i][j] = diagonals[i - j].pop(0)
        
        return mat

Complexity

  • ⏰ Time complexity: O((m + n) * d log(d)). Sorting each diagonal will take O(d log d), where d is the length of the diagonal. Since there are about m + n - 1 diagonals, the overall complexity is O((m + n) * min(m, n) log(min(m, n))).
  • 🧺 Space complexity: O(d), where the space is used for storing each diagonal’s elements.