1314. 矩阵区域和（二维前缀和）

package com.heu.wsq.leetcode.arr;

/**
 * 1314. 矩阵区域和
 * @author wsq
 * @date 2021/4/6
 * 给你一个 m * n 的矩阵 mat 和一个整数 K ，请你返回一个矩阵 answer ，其中每个 answer[i][j] 是所有满足下述条件的元素 mat[r][c] 的和： 
 * i - K <= r <= i + K, j - K <= c <= j + K 
 * (r, c) 在矩阵内。
 *
 * 示例 1：
 * 输入：mat = [[1,2,3],[4,5,6],[7,8,9]], K = 1
 * 输出：[[12,21,16],[27,45,33],[24,39,28]]
 *
 * 示例 2：
 * 输入：mat = [[1,2,3],[4,5,6],[7,8,9]], K = 2
 * 输出：[[45,45,45],[45,45,45],[45,45,45]]
 *
 * 链接：https://leetcode-cn.com/problems/matrix-block-sum
 */
public class MatrixBlockSum {
    public int[][] matrixBlockSum(int[][] mat, int k) {
        int m = mat.length;
        int n = mat[0].length;
        // 状态数组要比原数组多一行一列，类似序列动态规划，前0行或0列所构成的前缀后都为0（后续计算统一化）
        int[][] f = new int[m + 1][n + 1];
        // 先初始化二维前缀和，
        // 二维前缀和f[i][j]表示从mat[0][0]到mat[i][j]所构成的前缀和
        // 通过图形表示，f[i][j]等价于 mat[i-1][j-1] + f[i-1][j] + f[i][j-1] - f[i-1][j-1]
        for(int i = 1; i <= m; i++){
            for(int j = 1; j <= n; j++){
                f[i][j] = mat[i-1][j-1] + f[i-1][j] + f[i][j-1] - f[i-1][j-1];
            }
        }
        // 构建最后结果数组
        int[][] ans = new int[m][n];
        // i - K <= r <= i + K, j - K <= c <= j + K，表示ans[i][j]保存的是mat[i-k][i+k]到mat[i+k][j+k]的前缀和
        for(int i = 0; i < m; i++){
            for(int j = 0; j < n; j++){
                int x1 = Math.max(0, i - k);
                int x2 = Math.min(m - 1, i + k);
                int y1 = Math.max(0, j - k);
                int y2 = Math.min(n - 1, j + k);
                ans[i][j] = f[x2 + 1][y2 + 1] - f[x1][y2 + 1] - f[x2 + 1][y1] + f[x1][y1];
            }
        }
        return ans;
    }
}