Problem statement

https://leetcode.com/problems/matrix-block-sum/

Solution

The idea is to use accumulate 2d sums. Then iterate over all elements and calculate sums.

Complexity

It is O(mn) for time and space.

Code

class Solution:
    def matrixBlockSum(self, mat, k): 
        m, n = len(mat), len(mat[0])
        dp = [[0] * (n+1) for _ in range(m+1)]
        ans = [[0] * n for _ in range(m)]

        for c, r in product(range(n), range(m)):
            dp[r+1][c+1] = dp[r+1][c] + dp[r][c+1] - dp[r][c] + mat[r][c]
            
        for c, r in product(range(n), range(m)):
            r1, r2 = max(0, r-k), min(m-1, r+k)
            c1, c2 = max(0, c-k), min(n-1, c+k)
            ans[r][c] = dp[r2+1][c2+1] - dp[r1][c2+1] - dp[r2+1][c1] + dp[r1][c1]
                
        return ans