矩阵区域和（LeetCode）

题目

给你一个 m x n 的矩阵 mat 和一个整数 k ，请你返回一个矩阵 answer ，其中每个 answer[i][j] 是所有满足下述条件的元素 mat[r][c] 的和：

• i - k <= r <= i + k,
• j - k <= c <= j + k 且
• (r, c) 在矩阵内。

解题

import numpy as np


def matrixBlockSum(mat, k):
    # 获取矩阵的大小
    m, n = len(mat), len(mat[0])

    # 初始化前缀和矩阵，大小为 (m + 1) x (n + 1)
    prefix_sum = np.zeros((m + 1, n + 1), dtype=int)

    # 构建前缀和矩阵
    for i in range(m):
        for j in range(n):
            prefix_sum[i + 1][j + 1] = mat[i][j] + prefix_sum[i][j + 1] + prefix_sum[i + 1][j] - prefix_sum[i][j]

    # 初始化结果矩阵
    answer = np.zeros((m, n), dtype=int)

    # 计算结果矩阵
    for i in range(m):
        for j in range(n):
            # 定义子矩阵的上下边界和左右边界
            r1 = max(0, i - k)
            c1 = max(0, j - k)
            r2 = min(m, i + k + 1)
            c2 = min(n, j + k + 1)

            # 使用前缀和计算子矩阵元素和
            answer[i][j] = prefix_sum[r2][c2] - prefix_sum[r1][c2] - prefix_sum[r2][c1] + prefix_sum[r1][c1]

    return answer.tolist()


# 示例使用
mat = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
]
k = 1
result = matrixBlockSum(mat, k)
print(result)

[[12, 21, 16], [27, 45, 33], [24, 39, 28]]