Sort the Matrix Diagonally

medium array matrix sorting

Problem

You are given an m × n integer matrix. A diagonal here means a line of cells running from the top-left toward the bottom-right: it begins at some cell in the top row or the leftmost column and keeps stepping one row down and one column right until it leaves the grid. Sort the values along every such diagonal in ascending order and return the resulting matrix.

Constraints: 1 ≤ m, n ≤ 100 and 1 ≤ mat[i][j] ≤ 100.

Inputmat = [[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]]

The diagonal through (0,0), (1,1), (2,2) holds [3, 2, 1] and becomes [1, 2, 3]; every other ↘ diagonal is sorted the same way.

matrix (rows separated by ";", values by ",")

def diagonal_sort(mat):
    m, n = len(mat), len(mat[0])

    def sort_diag(i, j):
        vals = []
        r, c = i, j
        while r < m and c < n:       # collect the diagonal
            vals.append(mat[r][c])
            r, c = r + 1, c + 1
        vals.sort()                  # sort its values
        r, c = i, j
        for v in vals:               # write them back
            mat[r][c] = v
            r, c = r + 1, c + 1

    for j in range(n):               # diagonals starting on row 0
        sort_diag(0, j)
    for i in range(1, m):            # diagonals starting on column 0
        sort_diag(i, 0)
    return mat

function diagonalSort(mat) {
  const m = mat.length, n = mat[0].length;
  const sortDiag = (i, j) => {
    const vals = [];
    for (let r = i, c = j; r < m && c < n; r++, c++)
      vals.push(mat[r][c]);          // collect the diagonal
    vals.sort((a, b) => a - b);      // sort its values
    let k = 0;
    for (let r = i, c = j; r < m && c < n; r++, c++)
      mat[r][c] = vals[k++];         // write them back
  };
  for (let j = 0; j < n; j++) sortDiag(0, j);
  for (let i = 1; i < m; i++) sortDiag(i, 0);
  return mat;
}

int[][] diagonalSort(int[][] mat) {
    int m = mat.length, n = mat[0].length;
    for (int j = 0; j < n; j++) sortDiag(mat, 0, j);
    for (int i = 1; i < m; i++) sortDiag(mat, i, 0);
    return mat;
}

void sortDiag(int[][] mat, int i, int j) {
    int m = mat.length, n = mat[0].length;
    List<Integer> vals = new ArrayList<>();
    for (int r = i, c = j; r < m && c < n; r++, c++)
        vals.add(mat[r][c]);               // collect the diagonal
    Collections.sort(vals);                // sort its values
    int k = 0;
    for (int r = i, c = j; r < m && c < n; r++, c++)
        mat[r][c] = vals.get(k++);         // write them back
}

vector<vector<int>> diagonalSort(vector<vector<int>>& mat) {
    int m = mat.size(), n = mat[0].size();
    auto sortDiag = [&](int i, int j) {
        vector<int> vals;
        for (int r = i, c = j; r < m && c < n; r++, c++)
            vals.push_back(mat[r][c]);     // collect the diagonal
        sort(vals.begin(), vals.end());    // sort its values
        int k = 0;
        for (int r = i, c = j; r < m && c < n; r++, c++)
            mat[r][c] = vals[k++];         // write them back
    };
    for (int j = 0; j < n; j++) sortDiag(0, j);
    for (int i = 1; i < m; i++) sortDiag(i, 0);
    return mat;
}

Explanation

The key insight is that every cell on the same ↘ diagonal shares a constant value of i − j. Step one row down and one column right and both i and j grow by one, so their difference never changes. That single observation splits the matrix into m + n − 1 completely independent groups, and "sort the matrix diagonally" becomes "sort each small group on its own".

So the algorithm is three moves per diagonal: collect the values by walking down-right from the diagonal's start cell, sort that little list ascending, and write the sorted values back along the very same walk. Because diagonals never share a cell, no write can disturb another diagonal. Every diagonal starts either in the top row or in the leftmost column, which is exactly how the code enumerates them.

Walking the default example [[3,3,1,1],[2,2,1,2],[1,1,1,2]]: the diagonal starting at (0,0) collects [3, 2, 1] and writes back [1, 2, 3]; the one at (0,1) collects [3, 1, 2] and writes [1, 2, 3]; the one at (0,2) holds [1, 2], already sorted; the diagonal at (1,0) turns [2, 1] into [1, 2]; the two corner diagonals are single cells. The result is [[1,1,1,1],[1,2,2,2],[1,2,3,3]].

An equivalent implementation groups values into a hash map keyed by i − j in one sweep, sorts each bucket, and refills the matrix in a second sweep. It is the same idea with the same complexity; the per-diagonal walk shown here just avoids the extra map.

For complexity, every cell belongs to exactly one diagonal, and a diagonal holds at most min(m, n) cells. Sorting all the buckets therefore costs at most O(m·n · log(min(m, n))) overall, and the scratch list for one diagonal needs only O(min(m, n)) extra space.

Time: O(m·n · log(min(m, n))) Space: O(min(m, n))