Problem statement

https://binarysearch.com/problems/Matrix-Rectangular-Sums/

Solution

Equal to Leetcode 1314. Matrix Block Sum.

Complexity

It is O(mn) for time and space.

Code

class Solution:
    def solve(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