Skip to content

1314. Matrix Block Sum

Difficulty Topics

Description

Given a m x 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

Solution

matrix-block-sum.py
class Solution:
    def matrixBlockSum(self, mat: List[List[int]], k: int) -> List[List[int]]:
        rows, cols = len(mat), len(mat[0])
        prefix = [[0] * (cols + 1) for _ in range(rows + 1)]

        for i in range(rows):
            for j in range(cols):
                prefix[i + 1][j + 1] = prefix[i + 1][j] + prefix[i][j + 1] - prefix[i][j] + mat[i][j]

        for i in range(rows):
            for j in range(cols):
                r1, r2 = max(0, i - k), min(rows, i + k + 1)
                c1, c2 = max(0, j - k), min(cols, j + k + 1)

                mat[i][j] = (prefix[r2][c2] - prefix[r1][c2]) - (prefix[r2][c1] - prefix[r1][c1])

        return mat