Sort the Matrix Diagonally
MedPrompt
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.
Example 1:
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]]
Constraints:
m == mat.lengthn == mat[i].length1 <= m, n <= 1001 <= mat[i][j] <= 100
Approaches
3 approaches with complexity analysis and trade-offs.
This approach leverages a key property of diagonals: all cells (r, c) on the same diagonal share the same value of r - c. We can use this value as a key in a HashMap to group all elements belonging to the same diagonal. By using a PriorityQueue as the value in the map, the elements for each diagonal are automatically sorted as they are inserted. Finally, we rebuild the matrix by polling the sorted elements from the priority queues.
Algorithm
- Create a
HashMapwhere the key is the diagonal identifier (r - c) and the value is a data structure to hold the diagonal's elements. - A
PriorityQueueis a good choice for the value, as it will automatically keep the elements in sorted order. - Iterate through every cell
(r, c)of the input matrixmat. - For each cell, calculate the diagonal key
d = r - c. - Add the element
mat[r][c]to thePriorityQueueassociated with the keyd. - After populating the
HashMapwith all elements, iterate through the matrixmatagain. - For each cell
(r, c), retrieve the smallest element from thePriorityQueuefor the diagonalr - c(usingpoll()) and place it back intomat[r][c]. - Return the modified matrix.
Walkthrough
The core idea is to map each diagonal to a unique identifier and then sort the elements for each identifier. The difference between the row and column index, r - c, is constant for all cells on a given diagonal, making it a perfect identifier.
-
Group Elements: We traverse the entire
m x nmatrix. For each cellmat[r][c], we compute its diagonal keyr - c. We use aHashMap<Integer, PriorityQueue<Integer>>to store the elements. The key is the diagonal identifier, and thePriorityQueuestores all elements from that diagonal, naturally keeping them in ascending order. -
Reconstruct Matrix: After the map is fully populated, we traverse the matrix a second time. For each cell
mat[r][c], we again compute its keyr - c, access the correspondingPriorityQueuein our map, andpoll()the smallest element (which is the next in the sorted sequence for that diagonal) to place it inmat[r][c].
This method is conceptually straightforward as it abstracts away the manual traversal of each diagonal path.
import java.util.HashMap;import java.util.Map;import java.util.PriorityQueue; class Solution { public int[][] diagonalSort(int[][] mat) { int m = mat.length; int n = mat[0].length; Map<Integer, PriorityQueue<Integer>> diagonals = new HashMap<>(); // Step 1: Group elements by diagonal key and store in PriorityQueues for (int r = 0; r < m; r++) { for (int c = 0; c < n; c++) { int diagonalKey = r - c; diagonals.putIfAbsent(diagonalKey, new PriorityQueue<>()); diagonals.get(diagonalKey).add(mat[r][c]); } } // Step 2: Reconstruct the matrix with sorted elements for (int r = 0; r < m; r++) { for (int c = 0; c < n; c++) { mat[r][c] = diagonals.get(r - c).poll(); } } return mat; }}Complexity
Time
O(m * n * log(L)), where L = min(m, n) We iterate through all `m * n` cells twice. In the first pass, each element is added to a `PriorityQueue`. The maximum size of any priority queue is `L`, so each insertion takes `O(log L)` time. The total time is dominated by this step.
Space
O(m * n) The `HashMap` needs to store all `m * n` elements from the original matrix.
Trade-offs
Pros
Relatively simple to implement.
Avoids complex logic for manually iterating through each diagonal's path.
Cons
Has the highest space complexity, as it requires storing all
m * nelements in memory.The time complexity is not optimal due to the logarithmic factor from the priority queue operations.
This approach directly simulates the process described in the problem. It iterates through each diagonal one by one, extracts its elements into a temporary list, sorts that list, and then places the sorted elements back onto the diagonal in the matrix. Diagonals are identified by their starting points, which are all located on the top row and the first column of the matrix.
Algorithm
- Identify the starting cell of each diagonal. All diagonals start at either the first row (
(0, c)) or the first column ((r, 0)). - Iterate through all starting cells:
(0, 0)to(0, n-1)and(1, 0)to(m-1, 0). - For each starting cell, perform the following steps:
- Create a temporary
ArrayListto store the elements of the current diagonal. - Traverse the diagonal from its starting cell
(r, c)by repeatedly moving to(r+1, c+1)until you go out of the matrix bounds. Add each element to the list. - Sort the temporary list using a standard comparison-based sort (e.g.,
Collections.sort()). - Traverse the diagonal again from the start.
- Place the sorted elements from the list back into the matrix cells along the diagonal.
- Create a temporary
Walkthrough
Instead of grouping all diagonals at once, this method processes them sequentially. The main challenge is to correctly iterate over every diagonal.
-
Identify Diagonal Starts: All
m + n - 1diagonals begin atmat[0][c]forcfrom0ton-1, or atmat[r][0]forrfrom1tom-1. We can loop through these starting coordinates. -
Process Each Diagonal: For each starting coordinate, we have a helper function that does the work: a. Extraction: It traverses the diagonal path starting from
(r, c)and moving to(r+1, c+1), adding eachmat[i][j]to a temporary list. b. Sorting: It sorts this temporary list.Collections.sort()in Java works well here. c. Placement: It traverses the same diagonal path again, this time updating the matrix cells with values from the sorted list in order.
This approach is more space-efficient than the HashMap approach because it only needs to store the elements of one diagonal at a time.
import java.util.ArrayList;import java.util.Collections;import java.util.List; class Solution { public int[][] diagonalSort(int[][] mat) { int m = mat.length; int n = mat[0].length; // Sort diagonals starting from the first row for (int c = 0; c < n; c++) { sortSingleDiagonal(mat, 0, c, m, n); } // Sort diagonals starting from the first column (skip (0,0) as it's done) for (int r = 1; r < m; r++) { sortSingleDiagonal(mat, r, 0, m, n); } return mat; } private void sortSingleDiagonal(int[][] mat, int startR, int startC, int m, int n) { List<Integer> diagonalElements = new ArrayList<>(); int r = startR; int c = startC; // Extract elements from the diagonal while (r < m && c < n) { diagonalElements.add(mat[r][c]); r++; c++; } // Sort the elements Collections.sort(diagonalElements); // Place sorted elements back into the matrix r = startR; c = startC; int index = 0; while (r < m && c < n) { mat[r][c] = diagonalElements.get(index++); r++; c++; } }}Complexity
Time
O((m+n) * L * log(L)), where L = min(m, n) There are `m + n - 1` diagonals. For each diagonal of length `k`, we perform a sort which takes `O(k log k)`. The total time is the sum of `k * log(k)` over all diagonals. A loose upper bound is `O(m * n * log(L))`.
Space
O(min(m, n)) The space is determined by the size of the temporary list used to sort a single diagonal. The longest possible diagonal has `min(m, n)` elements.
Trade-offs
Pros
Excellent space complexity, as it only requires space for the longest diagonal.
Cons
The time complexity is not optimal due to using a comparison-based sort (
O(k log k)) for each diagonal.The implementation is slightly more involved as it requires logic to explicitly iterate along each diagonal path.
This is the most optimal approach, which refines the previous method by taking advantage of the problem's constraint: 1 <= mat[i][j] <= 100. Since the range of values is small and fixed, we can sort each diagonal in linear time using Counting Sort instead of a slower comparison-based sort. This eliminates the logarithmic factor from the time complexity, leading to a linear time solution overall.
Algorithm
- This approach builds upon the previous one (Iterate Diagonals and Sort) but replaces the comparison sort with a more efficient sorting algorithm.
- Given the constraint that matrix values are between 1 and 100, Counting Sort is an ideal choice.
- The overall structure remains the same: iterate through the starting cell of each diagonal.
- For each diagonal:
- Instead of a list, create a frequency array (or counting array) of size 101, initialized to zeros. Let's call it
counts. - Traverse the diagonal, and for each element
val, increment its frequency:counts[val]++. - After counting, traverse the diagonal again to place the sorted elements back.
- Keep a pointer
currentValto the current value to be placed, starting at 1. - For each cell on the diagonal, find the smallest
currentValfor whichcounts[currentVal]is greater than 0. PlacecurrentValin the cell, decrementcounts[currentVal], and then move to the next cell.
- Instead of a list, create a frequency array (or counting array) of size 101, initialized to zeros. Let's call it
Walkthrough
The logic is identical to the 'Iterate Diagonals and Sort' approach, but the sorting implementation within the helper function is replaced with Counting Sort for a significant performance boost.
-
Iteration: We still iterate through the starting points of each diagonal on the first row and first column.
-
Optimized Sorting: The
sortSingleDiagonalhelper is modified: a. Counting: Instead of adding elements to a list, we create a counting array of size 101 (for values 1-100). We traverse the diagonal and for each elementval, we incrementcounts[val]. b. Placement: We traverse the diagonal a second time to overwrite the cells. We use a pointer,currentVal, initialized to 1. For each cell, we find the next available number by advancingcurrentValuntilcounts[currentVal] > 0. We then placecurrentValinto the matrix cell, decrement its count in thecountsarray, and proceed to the next cell. If the count forcurrentValbecomes zero, the next cell will again search for the next available number.
This change reduces the sorting time for a diagonal of length k from O(k log k) to O(k + C), where C is the range of values (100). Since C is a constant, this is effectively O(k). Summing over all cells results in a linear time complexity for the entire matrix.
class Solution { public int[][] diagonalSort(int[][] mat) { int m = mat.length; int n = mat[0].length; // Sort diagonals starting from the first row for (int c = 0; c < n; c++) { sortDiagonalWithCountingSort(mat, 0, c, m, n); } // Sort diagonals starting from the first column for (int r = 1; r < m; r++) { sortDiagonalWithCountingSort(mat, r, 0, m, n); } return mat; } private void sortDiagonalWithCountingSort(int[][] mat, int startR, int startC, int m, int n) { // Values are 1 <= mat[i][j] <= 100 int[] counts = new int[101]; int r = startR; int c = startC; // Step 1: Count frequencies of numbers in the diagonal while (r < m && c < n) { counts[mat[r][c]]++; r++; c++; } // Step 2: Place numbers back in sorted order r = startR; c = startC; int currentVal = 1; while (r < m && c < n) { // Find the next smallest number available while (counts[currentVal] == 0) { currentVal++; } mat[r][c] = currentVal; counts[currentVal]--; r++; c++; } }}Complexity
Time
O(m * n) Each cell in the matrix is visited exactly twice: once to count its value and once to have a new value written to it. The work done per diagonal of length `k` is `O(k + C)` where `C` is the range of values (100). Summing over all diagonals, the total time is proportional to the total number of cells, `m * n`.
Space
O(1) The extra space is for the `counts` array. Since the range of values (1-100) is fixed by the problem constraints, the size of this array is constant (101) and does not depend on the input size `m` or `n`.
Trade-offs
Pros
Optimal time complexity of O(m * n).
Optimal space complexity of O(1) (since the counting array size is constant).
Cons
This approach is specialized and only works efficiently because of the small, fixed range of integer values in the matrix.
Solutions
Solution
public class Solution { public int[][] DiagonalSort(int[][] mat) { int m = mat.Length; int n = mat[0].Length; List < List < int >> g = new List < List < int >> (); for (int i = 0; i < m + n; i++) { g.Add(new List < int > ()); } for (int i = 0; i < m; i++) { for (int j = 0; j < n; j++) { g[m - i + j].Add(mat[i][j]); } } foreach(var e in g) { e.Sort((a, b) => b.CompareTo(a)); } for (int i = 0; i < m; i++) { for (int j = 0; j < n; j++) { int val = g[m - i + j][g[m - i + j].Count - 1]; g[m - i + j].RemoveAt(g[m - i + j].Count - 1); mat[i][j] = val; } } return mat; }}class Solution {public int[][] diagonalSort(int[][] mat) { int m = mat.length, n = mat[0].length; for (int k = 0; k < Math.min(m, n) - 1; ++k) { for (int i = 0; i < m - 1; ++i) { for (int j = 0; j < n - 1; ++j) { if (mat[i][j] > mat[i + 1][j + 1]) { int t = mat[i][j]; mat[i][j] = mat[i + 1][j + 1]; mat[i + 1][j + 1] = t; } } } } return mat; }}class Solution {public: vector<vector<int>> diagonalSort(vector<vector<int>> &mat) { int m = mat.size(), n = mat[0].size(); for (int k = 0; k < min(m, n) - 1; ++k) for (int i = 0; i < m - 1; ++i) for (int j = 0; j < n - 1; ++j) if (mat[i][j] > mat[i + 1][j + 1]) swap(mat[i][j], mat[i + 1][j + 1]); return mat; }};class Solution: def diagonalSort(self, mat: List[List[int]]) -> List[List[int]]: m, n = len(mat), len(mat[0]) for k in range(min(m, n) - 1): for i in range(m - 1): for j in range(n - 1): if mat[i][j] > mat[i + 1][j + 1]: mat[i][j], mat[i + 1][j + 1] = mat[i + 1][j + 1], mat[i][j] return matVideo walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.