Matrix Block Sum LeetCode solution in Rust

1
impl Solution {
2
  pub fn matrix_block_sum(mat: Vec<Vec<i32>>, k: i32) -> Vec<Vec<i32>> {
3
    let n = mat.len();
4
    let m = mat[0].len();
5
    let k = k as usize;
6
    let mut memo = vec![vec![0; m + 1]; n + 1];
7
    for i in 0..n {
8
      for j in 0..m {
9
        memo[i + 1][j + 1] = memo[i + 1][j] + memo[i][j + 1] - memo[i][j] + mat[i][j];
10
      }
11
    }
12

13
    let mut result = vec![vec![0; m]; n];
14
    for i in 0..n {
15
      for j in 0..m {
16
        let i_min = i.saturating_sub(k);
17
        let i_max = std::cmp::min(i + k, n - 1);
18
        let j_min = j.saturating_sub(k);
19
        let j_max = std::cmp::min(j + k, m - 1);
20

21
        result[i][j] = memo[i_max + 1][j_max + 1] - memo[i_max + 1][j_min] - memo[i_min][j_max + 1]
22
          + memo[i_min][j_min];
23
      }
24
    }
25
    result
26
  }
27
}