Code Implementation

class Solution:
    def diagonalSort(self, mat):
        from collections import defaultdict
        m, n = len(mat), len(mat[0])
        diags = defaultdict(list)
        # Collect elements of each diagonal
        for i in range(m):
            for j in range(n):
                diags[i - j].append(mat[i][j])
        # Sort each diagonal
        for key in diags:
            diags[key].sort(reverse=True)
        # Place sorted elements back
        for i in range(m):
            for j in range(n):
                mat[i][j] = diags[i - j].pop()
        return mat
      

class Solution {
public:
    vector<vector<int>> diagonalSort(vector<vector<int>>& mat) {
        int m = mat.size(), n = mat[0].size();
        unordered_map<int, vector<int>> diags;
        // Collect elements of each diagonal
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                diags[i - j].push_back(mat[i][j]);
            }
        }
        // Sort each diagonal
        for (auto &p : diags) {
            sort(p.second.begin(), p.second.end(), greater<int>());
        }
        // Place sorted elements back
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                mat[i][j] = diags[i - j].back();
                diags[i - j].pop_back();
            }
        }
        return mat;
    }
};
      

class Solution {
    public int[][] diagonalSort(int[][] mat) {
        int m = mat.length, n = mat[0].length;
        Map<Integer, List<Integer>> diags = new HashMap<>();
        // Collect elements of each diagonal
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                diags.computeIfAbsent(i - j, k -> new ArrayList<>()).add(mat[i][j]);
            }
        }
        // Sort each diagonal
        for (List<Integer> diag : diags.values()) {
            diag.sort(Comparator.reverseOrder());
        }
        // Place sorted elements back
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                mat[i][j] = diags.get(i - j).remove(diags.get(i - j).size() - 1);
            }
        }
        return mat;
    }
}
      

var diagonalSort = function(mat) {
    const m = mat.length, n = mat[0].length;
    const diags = {};
    // Collect elements of each diagonal
    for (let i = 0; i < m; ++i) {
        for (let j = 0; j < n; ++j) {
            const key = i - j;
            if (!(key in diags)) diags[key] = [];
            diags[key].push(mat[i][j]);
        }
    }
    // Sort each diagonal
    for (const key in diags) {
        diags[key].sort((a, b) => b - a);
    }
    // Place sorted elements back
    for (let i = 0; i < m; ++i) {
        for (let j = 0; j < n; ++j) {
            const key = i - j;
            mat[i][j] = diags[key].pop();
        }
    }
    return mat;
};
      

Problem Description

Given a 2D integer matrix mat of size m x n, your task is to sort each diagonal of the matrix in ascending order, where a diagonal runs from the top-left to the bottom-right. After sorting, the matrix should be updated in-place so that each diagonal contains its sorted values. Each element must remain within its original diagonal, and every diagonal must be sorted independently. There is only one valid solution for each input, and you must not reuse or move elements between different diagonals.

Thought Process

The core challenge is to sort elements that are connected diagonally, not row-wise or column-wise. At first glance, a brute-force approach might try to extract each diagonal, sort it, and put it back. However, we need an efficient way to map each cell to its diagonal and process all diagonals systematically.

Thinking further, we realize all elements on the same diagonal share the same difference between their row and column indices (i - j). This makes it easy to group elements by diagonal. By leveraging this property, we can collect, sort, and redistribute diagonal elements with minimal overhead.

The main optimization is to use a hash map (dictionary) to group elements by their diagonal key, sort them, and then write them back to their original positions.

Solution Approach

To solve the problem efficiently, we follow these steps:

1. Group elements by diagonal:
   • For each cell (i, j) in the matrix, compute the diagonal key as i - j.
   • Store the value mat[i][j] in a hash map or dictionary under the key i - j.
2. Sort each diagonal:
   • For each diagonal in the hash map, sort the list of values in ascending order.
   • To make writing back efficient, you can sort in reverse and use pop() from the list for fast removal of the smallest element.
3. Write back sorted values:
   • Iterate again through the matrix, and for each cell (i, j), assign it the next sorted value from its diagonal's list.
   • Ensure that each diagonal's values are used exactly once and in order.

This approach ensures that each diagonal is sorted independently and that all elements remain within their original diagonals.

Example Walkthrough

Consider the input matrix:

[[3, 3, 1, 1],
[2, 2, 1, 2],
[1, 1, 1, 2]]

1. Grouping by diagonal:
   • Diagonal 0 (i - j = 0): [3, 2, 1]
   • Diagonal -1: [3, 1, 2]
   • Diagonal -2: [1, 1]
   • Diagonal -3: [1]
   • Diagonal 1: [2, 1]
   • Diagonal 2: [1]
2. Sorting diagonals:
   • Diagonal 0: [1, 2, 3]
   • Diagonal -1: [1, 2, 3]
   • Diagonal -2: [1, 1]
   • Diagonal -3: [1]
   • Diagonal 1: [1, 2]
   • Diagonal 2: [1]
3. Writing back:
   • Fill each cell with the next sorted value from its diagonal.

The final sorted matrix is:
[[1, 1, 1, 1],
[1, 2, 2, 2],
[1, 2, 3, 3]]

Each diagonal is now sorted in ascending order, and elements remain within their original diagonals.

Time and Space Complexity

Brute-force approach:

• Extracting and sorting each diagonal individually without using a hash map would require scanning the matrix multiple times, leading to higher time complexity.
• Time complexity: O(mn log k), where k is the length of the largest diagonal.
• Space complexity: O(mn), since we need to store all elements somewhere while sorting.

• We visit each element exactly three times (grouping, sorting, writing back).
• Time complexity: O(mn log k), because sorting each diagonal of length k takes O(k log k), and the sum of all diagonal lengths is O(mn).
• Space complexity: O(mn), for storing all diagonal values in the hash map.

Summary

The key insight is that all cells on the same diagonal share the same i - j value, allowing us to group and process diagonals efficiently. By collecting elements into a hash map, sorting, and writing them back, we achieve an elegant and efficient solution that works for any matrix size. This approach avoids unnecessary complexity and leverages simple data structures for clarity and speed.