leetcode 1314 矩阵区域和

题意简述

给出一个矩阵，求以每个位置为中心，边长为$2k+1$的正方形区域内，覆盖到的原矩阵元素的总和。

算法分析

二维前缀和，前缀和矩阵留出第$0$行和第$0$列，注意正方形左上出界时，下标$(x, y)$应取为$(max(x, 0), max(y, 0))$。

代码实现

class Solution {
public:
    vector<vector<int>> matrixBlockSum(vector<vector<int>>& mat, int k) {
        int n = mat.size(), m = mat[0].size();
        vector<vector<int>> sum(n + 1, vector<int>(m + 1, 0));
        vector<vector<int>> ans(n, vector<int>(m, 0));

        for (int i = 1; i <= n; i++) {
            for (int j = 1; j <= m; j++) {
                sum[i][j] = 
                    sum[i - 1][j] + 
                    sum[i][j - 1] - 
                    sum[i - 1][j - 1] + 
                    mat[i - 1][j - 1];
            }
        }

        auto get_sum = [&](int i, int j, int k) {
            return 
                sum[min(i + k, n)][min(j + k, m)] - 
                sum[max(0, i - k - 1)][min(j + k, m)] - 
                sum[min(i + k, n)][max(j - k - 1, 0)] + 
                sum[max(0, i - k - 1)][max(0, j - k - 1)]
            ;
        };
        
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) {
                ans[i][j] = get_sum(i + 1, j + 1, k);
            }
        }

        return ans;
    }
};