Matrix Block Sum LeetCode solution in C++

1
class Solution {
2
 public:
3
  vector<vector<int>> matrixBlockSum(vector<vector<int>> &mat, int k) {
4
    int n = mat.size(), m = mat[0].size();
5
    vector<vector<int>> pref(n + 1, vector<int>(m + 1, 0));
6
    // Calculating prefix sum
7
    for (int i = 1; i <= n; i++) {
8
      for (int j = 1; j <= m; j++) {
9
        // note that while counting for [i-1][j] and [i][j-1];
10
        // pref[i-1][j-1] will be added twice. So we reduce it once.
11
        pref[i][j] = mat[i - 1][j - 1] + pref[i - 1][j] + pref[i][j - 1] - pref[i - 1][j - 1];
12
      }
13
    }
14

15
    // Find the sum of the elements specified in the K-block
16
    vector<vector<int>> res(n, vector<int>(m, 0));
17
    for (int i = 0; i < n; i++) {
18
      for (int j = 0; j < m; j++) {
19
        // checking for all pairs to be in bounds.
20
        int r1 = max(0, i - k);
21
        int c1 = max(0, j - k);
22
        int r2 = min(n - 1, i + k);
23
        int c2 = min(m - 1, j + k);
24
        // finding res[i][j] = bottom right - bottom left - top right + top left
25
        res[i][j] = pref[r2 + 1][c2 + 1] - pref[r2 + 1][c1] - pref[r1][c2 + 1] + pref[r1][c1];
26
      }
27
    }
28

29
    return res;
30
  }
31
};