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

    unordered_map<int, priority_queue<int>> count;

    for (int i = 0; i < m; ++i)
      for (int j = 0; j < n; ++j) count[i - j].push(mat[i][j]);

    for (int i = m - 1; i >= 0; --i)
      for (int j = n - 1; j >= 0; --j)
        mat[i][j] = count[i - j].top(), count[i - j].pop();

    return mat;
  }
};

class Solution {
  public int[][] diagonalSort(int[][] mat) {
    final int m = mat.length;
    final int n = mat[0].length;

    Map<Integer, Queue<Integer>> count = new HashMap<>();

    for (int i = 0; i < m; ++i)
      for (int j = 0; j < n; ++j)
        count.computeIfAbsent(i - j, k -> new PriorityQueue<>()).add(mat[i][j]);

    for (int i = 0; i < m; ++i)
      for (int j = 0; j < n; ++j)
        mat[i][j] = count.get(i - j).poll();

    return mat;
  }
}

class Solution:
  def diagonalSort(self, mat: List[List[int]]) -> List[List[int]]:
    m = len(mat)
    n = len(mat[0])

    count = defaultdict(list)

    for i in range(m):
      for j in range(n):
        count[i - j].append(mat[i][j])

    for value in count.values():
      value.sort(reverse=1)

    for i in range(m):
      for j in range(n):
        mat[i][j] = count[i - j].pop()

    return mat