1
impl Solution {
2
  pub fn diagonal_sort(mut mat: Vec<Vec<i32>>) -> Vec<Vec<i32>> {
3
    let mut buckets = [0; 101];
4
    let (m, n) = (mat.len(), mat[0].len());
5

6
    // Chain together iterators for first element coordinates of all relevant diagonals
7
    (1..m - 1)
8
      .map(|i| (i, 0))
9
      .chain((0..n - 1).map(|j| (0, j)))
10
      .for_each(|(i0, j0)| {
11
        let (mut i, mut j) = (i0, j0);
12
        let (mut min, mut max) = (i32::MAX, i32::MIN);
13
        // Populate buckets for bucket sorting
14
        while i < m && j < n {
15
          let element = mat[i][j];
16
          min = min.min(element);
17
          max = max.max(element);
18
          buckets[element as usize] += 1;
19
          i += 1;
20
          j += 1;
21
        }
22
        // Populate matrix with result of bucket sorting
23
        let (mut i, mut j) = (i0, j0);
24
        let mut k = min as usize;
25
        while i < m && j < n {
26
          while buckets[k] == 0 {
27
            k += 1;
28
          }
29
          mat[i][j] = k as i32;
30
          buckets[k] -= 1;
31
          i += 1;
32
          j += 1;
33
        }
34
      });
35
    mat
36
  }
37
}