Sort the Matrix Diagonally

Med
#1235Time: 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.2 companies
Algorithms
Data structures
Companies

Prompt

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.length
  • n == mat[i].length
  • 1 <= m, n <= 100
  • 1 <= 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 HashMap where the key is the diagonal identifier (r - c) and the value is a data structure to hold the diagonal's elements.
  • A PriorityQueue is 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 matrix mat.
  • For each cell, calculate the diagonal key d = r - c.
  • Add the element mat[r][c] to the PriorityQueue associated with the key d.
  • After populating the HashMap with all elements, iterate through the matrix mat again.
  • For each cell (r, c), retrieve the smallest element from the PriorityQueue for the diagonal r - c (using poll()) and place it back into mat[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.

  1. Group Elements: We traverse the entire m x n matrix. For each cell mat[r][c], we compute its diagonal key r - c. We use a HashMap<Integer, PriorityQueue<Integer>> to store the elements. The key is the diagonal identifier, and the PriorityQueue stores all elements from that diagonal, naturally keeping them in ascending order.

  2. 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 key r - c, access the corresponding PriorityQueue in our map, and poll() the smallest element (which is the next in the sorted sequence for that diagonal) to place it in mat[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 * n elements in memory.

  • The time complexity is not optimal due to the logarithmic factor from the priority queue operations.

Solutions

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;    }}

Video 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.