Sort Matrix by Diagonals
MedPrompt
You are given an n x n square matrix of integers grid. Return the matrix such that:
- The diagonals in the bottom-left triangle (including the middle diagonal) are sorted in non-increasing order.
- The diagonals in the top-right triangle are sorted in non-decreasing order.
Example 1:
Input: grid = [[1,7,3],[9,8,2],[4,5,6]]
Output: [[8,2,3],[9,6,7],[4,5,1]]
Explanation:
The diagonals with a black arrow (bottom-left triangle) should be sorted in non-increasing order:
[1, 8, 6]becomes[8, 6, 1].[9, 5]and[4]remain unchanged.
The diagonals with a blue arrow (top-right triangle) should be sorted in non-decreasing order:
[7, 2]becomes[2, 7].[3]remains unchanged.
Example 2:
Input: grid = [[0,1],[1,2]]
Output: [[2,1],[1,0]]
Explanation:
The diagonals with a black arrow must be non-increasing, so [0, 2] is changed to [2, 0]. The other diagonals are already in the correct order.
Example 3:
Input: grid = [[1]]
Output: [[1]]
Explanation:
Diagonals with exactly one element are already in order, so no changes are needed.
Constraints:
grid.length == grid[i].length == n1 <= n <= 10-105 <= grid[i][j] <= 105
Approaches
2 approaches with complexity analysis and trade-offs.
This approach uses a HashMap to group the elements of the matrix by their diagonals. The key for the map is an integer that uniquely identifies a diagonal, which can be calculated as i - j for an element at grid[i][j]. After grouping all elements, each diagonal's list of elements is sorted according to the problem's rules and then placed back into the matrix.
Algorithm
- Create a
HashMap<Integer, List<Integer>>to map diagonal indices (i-j) to lists of their elements. - Iterate through the
gridfrom(0,0)to(n-1, n-1). For each elementgrid[i][j], add it to the list for diagonali-j. - Iterate through the entries in the HashMap.
- If the corresponding key
kis>= 0, sort the list in descending (non-increasing) order. - If the key
kis< 0, sort the list in ascending (non-decreasing) order. - Create a
HashMap<Integer, Integer>to act as pointers for each sorted diagonal list, initialized to 0. - Iterate through the
gridagain. For eachgrid[i][j], get the element from the sorted list for diagonali-jat the current pointer index, place it ingrid[i][j], and increment the pointer for that diagonal. - Return the modified
grid.
Walkthrough
We first iterate through the entire n x n matrix. For each element grid[i][j], we compute its diagonal index k = i - j. We use a HashMap<Integer, List<Integer>> where the key is the diagonal index k and the value is a list of elements on that diagonal. We add grid[i][j] to the list corresponding to key k.
After populating the map, we iterate through each entry in the map. For each diagonal (list of elements), we sort it. If the diagonal index k is non-negative (k >= 0), we sort the list in non-increasing (descending) order. Otherwise (k < 0), we sort it in non-decreasing (ascending) order.
Since the elements were added to the lists by traversing the matrix row by row, the elements within each list are already ordered from top-left to bottom-right along their diagonal. To place them back correctly after sorting, we need to keep track of which element to pick from the sorted list for each diagonal. We can use another map, say pointers, to store the current index for each diagonal's sorted list.
Finally, we iterate through the matrix a second time. For each position (i, j), we find its diagonal index k = i - j, get the next available element from the sorted list for diagonal k using our pointers map, and place it in the result matrix at grid[i][j]. We then increment the pointer for diagonal k.
import java.util.*; class Solution { public int[][] diagonalSort(int[][] grid) { int n = grid.length; Map<Integer, List<Integer>> diagonals = new HashMap<>(); // 1. Store all diagonals in a HashMap for (int i = 0; i < n; i++) { for (int j = 0; j < n; j++) { int key = i - j; diagonals.putIfAbsent(key, new ArrayList<>()); diagonals.get(key).add(grid[i][j]); } } // 2. Sort each diagonal for (int key : diagonals.keySet()) { List<Integer> diag = diagonals.get(key); if (key >= 0) { // Bottom-left triangle and main diagonal: non-increasing Collections.sort(diag, Collections.reverseOrder()); } else { // Top-right triangle: non-decreasing Collections.sort(diag); } } // 3. Place sorted elements back into the grid Map<Integer, Integer> pointers = new HashMap<>(); for (int i = 0; i < n; i++) { for (int j = 0; j < n; j++) { int key = i - j; int index = pointers.getOrDefault(key, 0); grid[i][j] = diagonals.get(key).get(index); pointers.put(key, index + 1); } } return grid; }}Complexity
Time
O(n^2 * log n) - Populating the map takes `O(n^2)`. Sorting all diagonals takes `sum(d_k * log(d_k))` where `d_k` is the length of diagonal `k`. Since `sum(d_k) = n^2` and `d_k <= n`, this is bounded by `O(n^2 * log n)`. Placing elements back takes `O(n^2)`.
Space
O(n^2) - The HashMap stores all `n^2` elements of the grid. The pointers map stores `2n-1` entries.
Trade-offs
Pros
Conceptually simple and directly follows the definition of diagonals (
i-j).Easy to implement without complex loop conditions for diagonal traversal.
Cons
Requires extra space proportional to the size of the matrix to store all elements in the HashMap.
Involves multiple data structures (two HashMaps) and multiple passes over the matrix data, which can have higher overhead.
This approach improves on space complexity by processing one diagonal at a time. Instead of storing all diagonals simultaneously, we iterate through the starting points of each diagonal, extract its elements into a temporary list, sort the list, and then immediately place the sorted elements back onto their diagonal in the matrix. This avoids the need for a large HashMap to store the entire matrix.
Algorithm
- Handle bottom-left and main diagonals:
- Loop
ifrom0ton-1(representing the starting row in the first column). - For each
i, create a temporary listdiag. - Traverse the diagonal starting at
(i, 0): while rowrand columncare in bounds, addgrid[r][c]todiag, then incrementrandc. - Sort
diagin descending (non-increasing) order. - Traverse the diagonal starting at
(i, 0)again, updatinggrid[r][c]with elements from the sorteddiag.
- Loop
- Handle top-right diagonals:
- Loop
jfrom1ton-1(representing the starting column in the first row). - For each
j, create a temporary listdiag. - Traverse the diagonal starting at
(0, j): whilerandcare in bounds, addgrid[r][c]todiag, then incrementrandc. - Sort
diagin ascending (non-decreasing) order. - Traverse the diagonal starting at
(0, j)again, updatinggrid[r][c]with elements from the sorteddiag.
- Loop
- Return the modified
grid.
Walkthrough
The key idea is to identify the starting cell of each diagonal. Diagonals in an n x n matrix can be uniquely identified by their starting cell, which will be in the first row or the first column.
We can split the process into two parts:
- Bottom-left and main diagonals: These diagonals start on the first column. We can iterate with a
start_rowfrom0ton-1. The starting cell for each of these diagonals is(start_row, 0). - Top-right diagonals: These diagonals start on the first row (excluding the main diagonal's start at
(0,0)which is already covered). We can iterate with astart_colfrom1ton-1. The starting cell is(0, start_col).
For each starting cell:
- We traverse the diagonal from that starting cell, collecting all its elements into a temporary list.
- We sort this temporary list. For diagonals starting at
(start_row, 0), the diagonal indexi-jisstart_row - 0 = start_row >= 0, so we sort in non-increasing order. For diagonals starting at(0, start_col), the index is0 - start_col < 0, so we sort in non-decreasing order. - We traverse the same diagonal again, this time updating the matrix cells with the values from the sorted temporary list.
This process is repeated for all 2n - 1 diagonals, modifying the grid in place.
import java.util.*; class Solution { public int[][] diagonalSort(int[][] grid) { int n = grid.length; // Process diagonals in the bottom-left triangle and the main diagonal // These start from the first column (j=0) for (int i = 0; i < n; i++) { List<Integer> diagonal = new ArrayList<>(); int r = i, c = 0; while (r < n && c < n) { diagonal.add(grid[r][c]); r++; c++; } // Sort non-increasingly Collections.sort(diagonal, Collections.reverseOrder()); r = i; c = 0; int k = 0; while (r < n && c < n) { grid[r][c] = diagonal.get(k++); r++; c++; } } // Process diagonals in the top-right triangle // These start from the first row (i=0), skipping the main diagonal's start (0,0) for (int j = 1; j < n; j++) { List<Integer> diagonal = new ArrayList<>(); int r = 0, c = j; while (r < n && c < n) { diagonal.add(grid[r][c]); r++; c++; } // Sort non-decreasingly Collections.sort(diagonal); r = 0; c = j; int k = 0; while (r < n && c < n) { grid[r][c] = diagonal.get(k++); r++; c++; } } return grid; }}Complexity
Time
O(n^2 * log n) - We iterate through each element twice (once to read, once to write). The total number of elements is `n^2`. The sorting cost dominates. For each diagonal of length `d`, we spend `O(d * log d)`. Summing over all diagonals gives `O(n^2 * log n)`.
Space
O(n) - The maximum space used is for the temporary list to store the elements of the longest diagonal, which is the main diagonal of length `n`.
Trade-offs
Pros
More space-efficient than the HashMap approach, as it only requires space for one diagonal at a time.
Modifies the grid in-place, avoiding the need for a separate result matrix.
Cons
The logic for iterating through diagonals by their start points might be slightly more complex to write correctly compared to the HashMap approach.
Solutions
Solution
class Solution {public int[][] sortMatrix(int[][] grid) { int n = grid.length; for (int k = n - 2; k >= 0; --k) { int i = k, j = 0; List<Integer> t = new ArrayList<>(); while (i < n && j < n) { t.add(grid[i++][j++]); } Collections.sort(t); for (int x : t) { grid[--i][--j] = x; } } for (int k = n - 2; k > 0; --k) { int i = k, j = n - 1; List<Integer> t = new ArrayList<>(); while (i >= 0 && j >= 0) { t.add(grid[i--][j--]); } Collections.sort(t); for (int x : t) { grid[++i][++j] = x; } } return grid; }}class Solution {public: vector<vector<int>> sortMatrix(vector<vector<int>> &grid) { int n = grid.size(); for (int k = n - 2; k >= 0; --k) { int i = k, j = 0; vector<int> t; while (i < n && j < n) { t.push_back(grid[i++][j++]); } ranges ::sort(t); for (int x : t) { grid[--i][--j] = x; } } for (int k = n - 2; k > 0; --k) { int i = k, j = n - 1; vector<int> t; while (i >= 0 && j >= 0) { t.push_back(grid[i--][j--]); } ranges ::sort(t); for (int x : t) { grid[++i][++j] = x; } } return grid; }};class Solution: def sortMatrix(self, grid: List[List[int]]) -> List[List[int]]: n = len(grid) for k in range(n - 2, - 1, - 1): i, j = k, 0 t = [] while i < n and j < n: t . append(grid[i][j]) i += 1 j += 1 t . sort() i, j = k, 0 while i < n and j < n: grid[i][j] = t . pop() i += 1 j += 1 for k in range(n - 2, 0, - 1): i, j = k, n - 1 t = [] while i >= 0 and j >= 0: t . append(grid[i][j]) i -= 1 j -= 1 t . sort() i, j = k, n - 1 while i >= 0 and j >= 0: grid[i][j] = t . pop() i -= 1 j -= 1 return gridVideo 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.