Sort the Matrix Diagonally — Editorial

Opening the reading…

Intuitionthe diagonal identity that turns a two-dimensional task into independent one-dimensional tasks

A diagonal moves one row down and one column right at every step. Because both coordinates increase together, their difference stays unchanged: a cell at row i and column j belongs to the diagonal identified by i - j. Cells with the same difference are therefore exactly the cells that must be sorted together, while cells with different differences must never be mixed.

The solution has two passes over the matrix. First, place every value into the collection for its diagonal and sort each collection. Second, walk through the matrix in the same order and take the next sorted value from that cell's collection. The traversal order matters: along a fixed diagonal, row-major matrix order visits cells from top-left to bottom-right, so consuming values from the front recreates the required direction.

The invariant i - j identifies each diagonal.

Approach

Read m and n from the matrix and create one collection for every diagonal key i - j, because grouping by row or column would combine cells that do not belong to the same diagonal.

Visit every cell in row-major order and append its value to diag[i - j], because equal differences identify the same bottom-right diagonal and this order already matches the order in which that diagonal will later be filled.

Sort every diagonal collection in ascending order, because the final matrix must be ordered independently along each diagonal rather than globally across all cells.

Create a next-position counter for every diagonal and start it at zero, because each key needs its own read position and sharing one counter would consume values from unrelated diagonals.

Visit every cell again, compute its key, and assign the next value from that key's sorted collection, because the second traversal writes values back along each diagonal in top-left to bottom-right order.

Advance only that key's counter after the assignment, because the following cell on the same diagonal must receive the following sorted value while other diagonals must keep their own positions.

Return the modified matrix, whose output cells are required by the problem and therefore are not counted as extra working storage.

Complexitysorting cost depends on the lengths of the independent diagonals

MEASURE	BOUND	WHY
Time	O(mn log(min(m, n))) worst case	Every cell is inserted once and written once, contributing O(mn). Sorting contributes the sum of k log k over all diagonals; since each diagonal has at most min(m, n) cells and all diagonal lengths sum to mn, this is at most O(mn log(min(m, n))).
Space	O(mn) extra, worst case	The diagonal collections together store exactly one working copy of every input value, and the counters use one entry per diagonal. The returned matrix is required output and is excluded. The bound reaches O(mn) because the implementation keeps all diagonal values until the write-back pass.

Here m is the number of rows, n is the number of columns, and k is the length of one diagonal. The worst-case bound uses the longest possible diagonal, min(m, n).

Annotated solutionC++ · hash-map buckets · collect, sort, and consume

CPPThe complete two-pass solution: bucket by i - j, sort each bucket, then consume it during write-back.

#include <algorithm>
#include <unordered_map>
#include <vector>

using namespace std;

class Solution {
public:
    vector<vector<int>> diagonalSort(vector<vector<int>>& mat) {
        int m = mat.size();
        int n = mat[0].size();

        unordered_map<int, vector<int>> diag;

        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                diag[i - j].push_back(mat[i][j]);
            }
        }

        for (auto& entry : diag) {
            sort(entry.second.begin(), entry.second.end());
        }

        unordered_map<int, int> next;
        for (const auto& entry : diag) {
            next[entry.first] = 0;
        }

        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                int key = i - j;
                mat[i][j] = diag[key][next[key]];
                ++next[key];
            }
        }

        return mat;
    }
};

The two maps have different jobs. diag stores the sorted values for each diagonal, while next stores how many values from that diagonal have already been written. Keeping the read position separate from the values avoids erasing the bucket or repeatedly taking its first element, either of which would make later cells expensive or incorrect.

Sorting one diagonal at a timean O(min(m, n)) extra-space alternative

You can avoid storing every diagonal simultaneously by starting at each cell in the top row and each cell in the leftmost column, collecting one diagonal, sorting it, and writing it back immediately. Each diagonal is still processed independently, but its temporary vector is released before the next diagonal starts. This is a genuine space optimisation, not merely a different arrangement: time remains O(mn log(min(m, n))), while extra space drops to O(min(m, n)).

CPPThe lower-space alternative: process every diagonal from its top row or leftmost-column starting cell.

#include <algorithm>
#include <vector>

using namespace std;

class Solution {
    void sortFrom(vector<vector<int>>& mat, int startRow, int startCol) {
        int m = mat.size();
        int n = mat[0].size();
        vector<int> values;

        int row = startRow;
        int col = startCol;
        while (row < m && col < n) {
            values.push_back(mat[row][col]);
            ++row;
            ++col;
        }

        sort(values.begin(), values.end());

        row = startRow;
        col = startCol;
        int index = 0;
        while (row < m && col < n) {
            mat[row][col] = values[index++];
            ++row;
            ++col;
        }
    }

public:
    vector<vector<int>> diagonalSort(vector<vector<int>>& mat) {
        int m = mat.size();
        int n = mat[0].size();

        for (int col = 0; col < n; ++col) {
            sortFrom(mat, 0, col);
        }
        for (int row = 1; row < m; ++row) {
            sortFrom(mat, row, 0);
        }

        return mat;
    }
};

The starting cells are exactly the top row and the leftmost column, with the top-left cell included only once. Starting elsewhere would repeat a diagonal, while omitting either border would leave some cells untouched. The bucket-map version is usually easier to explain and less vulnerable to start-point mistakes; the border-walk version is useful when peak working memory matters.