【leetcode】1314. Matrix Block Sum

題目如下：

Given a m * n matrix mat and an integer K, return a matrix answer where each answer[i][j] is the sum of all elements mat[r][c] for i - K <= r <= i + K, j - K <= c <= j + K, and (r, c) is a valid position in the matrix.

Example 1:

Input: mat = [[1,2,3],[4,5,6],[7,8,9]], K = 1
Output: [[12,21,16],[27,45,33],[24,39,28]]

Example 2:

Input: mat = [[1,2,3],[4,5,6],[7,8,9]], K = 2
Output: [[45,45,45],[45,45,45],[45,45,45]]

Constraints:

• m == mat.length
• n == mat[i].length
• 1 <= m, n, K <= 100
• 1 <= mat[i][j] <= 100

解題思路：記grid[i][j] 為 左上角頂點是(0,0)，右下角頂點是(i,j) 的矩陣和，grid可以通過時間復雜度為O(n^2)的循環計算出來，接下來再遍歷mat，計算出相應每個子矩陣和即可。

代碼如下：

class Solution(object):
    def matrixBlockSum(self, mat, K):
        """
        :type mat: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        res = [[0] * len(mat[0]) for _ in mat]
        grid = [[0] * len(mat[0]) for _ in mat]
        for i in range(len(mat)):
            count = 0
            for j in range(len(mat[0])):
                count += mat[i][j]
                grid[i][j] = count

        for i in range(len(res)):
            for j in range(len(res[i])):
                left = max(0,j-K)
                right = min(j+K,len(res[i])-1)
                count = 0
                for row in range(max(0,i-K),min(i+K+1,len(res))):
                    if left == 0:
                        count += grid[row][right]
                    else:
                        count += (grid[row][right] - grid[row][left-1])
                res[i][j] = count
        return res