Problem Description

You are given a m x n matrix mat and an integer k. Your task is to create a new matrix answer where each element answer[i][j] represents the sum of all elements from the original matrix within a specific rectangular region centered around position (i, j).

For each position (i, j) in the result matrix, you need to sum all elements mat[r][c] from the original matrix where:

• The row r is within the range: i - k ≤ r ≤ i + k
• The column c is within the range: j - k ≤ c ≤ j + k
• The position (r, c) must be valid (within the bounds of the original matrix)

In other words, for each cell in the result matrix, you calculate the sum of all elements in a rectangular block that extends k positions in all four directions (up, down, left, right) from that cell's corresponding position in the original matrix. If the block extends beyond the matrix boundaries, only the valid portions within the matrix are included in the sum.

For example, if k = 1, then for position (i, j), you would sum all elements in a 3×3 block centered at (i, j) (or smaller if near the edges of the matrix).

Intuition

The naive approach would be to calculate the sum for each block independently. For each position (i, j), we'd iterate through all elements in its block and sum them up. However, this would be inefficient because we'd be recalculating overlapping regions multiple times.

The key insight is that we're repeatedly calculating sums of rectangular regions in the matrix. When we move from one position to the next, the blocks overlap significantly, meaning we're summing many of the same elements again. This suggests we need a way to reuse previous calculations.

This is where the concept of prefix sums comes in. In one dimension, if we know the cumulative sum up to each position, we can quickly calculate the sum of any range. For example, to get the sum from index a to b, we just compute prefix[b] - prefix[a-1].

We can extend this idea to two dimensions. If we precompute the cumulative sum for every rectangular region starting from the top-left corner (0, 0) to any position (i, j), we can then calculate the sum of any arbitrary rectangular region using inclusion-exclusion principle.

Think of it like this: to get the sum of a rectangle, we take the total sum up to the bottom-right corner, subtract the excess regions on the left and top, but then we need to add back the top-left corner area that we subtracted twice. This gives us the formula:

sum = s[x2+1][y2+1] - s[x1][y2+1] - s[x2+1][y1] + s[x1][y1]

By preprocessing these prefix sums once at the beginning, we can answer each block sum query in constant time O(1), making the overall solution much more efficient than the naive approach.

Learn more about Prefix Sum patterns.

Solution Approach

The solution uses a two-dimensional prefix sum technique to efficiently calculate block sums.

Step 1: Build the Prefix Sum Array

First, we create a 2D prefix sum array s with dimensions (m+1) x (n+1), where the extra row and column of zeros simplify boundary handling.

For each position (i, j) in the prefix sum array, s[i][j] represents the sum of all elements in the rectangle from (0, 0) to (i-1, j-1) in the original matrix.

The formula to build this array is:

s[i][j] = s[i-1][j] + s[i][j-1] - s[i-1][j-1] + mat[i-1][j-1]

This works because:

• s[i-1][j] gives us the sum of the rectangle above the current row
• s[i][j-1] gives us the sum of the rectangle to the left of the current column
• We subtract s[i-1][j-1] because it was counted twice (in both previous terms)
• We add the current element mat[i-1][j-1]

Step 2: Calculate Block Sums Using Prefix Sums

For each position (i, j) in the result matrix, we need to:

1. Determine the boundaries of the block:

   • Top-left corner: (x1, y1) = (max(i-k, 0), max(j-k, 0))
   • Bottom-right corner: (x2, y2) = (min(i+k, m-1), min(j+k, n-1))

   These max and min operations ensure we stay within matrix bounds.

2. Calculate the sum using the inclusion-exclusion principle:

   ans[i][j] = s[x2+1][y2+1] - s[x1][y2+1] - s[x2+1][y1] + s[x1][y1]

   This formula works by:

   • Starting with the total sum from (0,0) to (x2, y2): s[x2+1][y2+1]
   • Subtracting the rectangle to the left of our region: s[x1][y2+1]
   • Subtracting the rectangle above our region: s[x2+1][y1]
   • Adding back the top-left rectangle we subtracted twice: s[x1][y1]

Time and Space Complexity:

• Time: O(m*n) for both building the prefix sum array and calculating all block sums
• Space: O(m*n) for the prefix sum array and result matrix

This approach transforms what would be O(m*n*k²) time complexity with the naive method into O(m*n), making it much more efficient for large values of k.

Example Walkthrough

Let's walk through a small example to illustrate the solution approach.

Given:

• Matrix mat = [[1,2,3],[4,5,6],[7,8,9]] (3×3 matrix)
• k = 1

Step 1: Build the Prefix Sum Array

We create a 4×4 prefix sum array s (one size larger than the original matrix):

Initially, s is all zeros:

s = [[0,0,0,0],
     [0,0,0,0],
     [0,0,0,0],
     [0,0,0,0]]

Now we fill it using the formula s[i][j] = s[i-1][j] + s[i][j-1] - s[i-1][j-1] + mat[i-1][j-1]:

• s[1][1] = 0 + 0 - 0 + 1 = 1
• s[1][2] = 0 + 1 - 0 + 2 = 3
• s[1][3] = 0 + 3 - 0 + 3 = 6
• s[2][1] = 1 + 0 - 0 + 4 = 5
• s[2][2] = 5 + 3 - 1 + 5 = 12
• s[2][3] = 12 + 6 - 3 + 6 = 21
• s[3][1] = 5 + 0 - 0 + 7 = 12
• s[3][2] = 12 + 12 - 5 + 8 = 27
• s[3][3] = 27 + 21 - 12 + 9 = 45

Final prefix sum array:

s = [[0, 0, 0, 0],
     [0, 1, 3, 6],
     [0, 5,12,21],
     [0,12,27,45]]

Step 2: Calculate Block Sums

Now let's calculate answer[1][1] (the center element):

1. Determine boundaries for position (1,1) with k=1:

   • Top-left: (x1, y1) = (max(1-1, 0), max(1-1, 0)) = (0, 0)
   • Bottom-right: (x2, y2) = (min(1+1, 2), min(1+1, 2)) = (2, 2)

   This gives us a 3×3 block covering the entire matrix.

2. Apply the formula:

   answer[1][1] = s[2+1][2+1] - s[0][2+1] - s[2+1][0] + s[0][0]
                = s[3][3] - s[0][3] - s[3][0] + s[0][0]
                = 45 - 0 - 0 + 0
                = 45

   This correctly sums all elements: 1+2+3+4+5+6+7+8+9 = 45

Let's calculate answer[0][0] (top-left corner):

1. Determine boundaries for position (0,0) with k=1:

   • Top-left: (x1, y1) = (max(0-1, 0), max(0-1, 0)) = (0, 0)
   • Bottom-right: (x2, y2) = (min(0+1, 2), min(0+1, 2)) = (1, 1)

   This gives us a 2×2 block in the top-left corner.

2. Apply the formula:

   answer[0][0] = s[1+1][1+1] - s[0][1+1] - s[1+1][0] + s[0][0]
                = s[2][2] - s[0][2] - s[2][0] + s[0][0]
                = 12 - 0 - 0 + 0
                = 12

   This correctly sums the 2×2 block: 1+2+4+5 = 12

The complete result matrix would be:

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

Each value represents the sum of all elements within distance k=1 from that position in the original matrix.

Solution Implementation

Time and Space Complexity

Time Complexity: O(m × n)

The algorithm consists of two main phases:

1. Building the prefix sum matrix: The nested loops iterate through all m rows and n columns of the input matrix exactly once. Each cell operation (computing s[i][j]) takes O(1) time, resulting in O(m × n) time.

2. Computing the answer matrix: Another set of nested loops iterates through all m × n positions. For each position (i, j), the algorithm:

   • Calculates boundaries (x1, y1, x2, y2) in O(1) time
   • Computes the block sum using the prefix sum formula in O(1) time

   This phase also takes O(m × n) time.

Total time complexity: O(m × n) + O(m × n) = O(m × n)

Space Complexity: O(m × n)

The algorithm uses:

• Prefix sum matrix s of size (m + 1) × (n + 1): O(m × n) space
• Answer matrix ans of size m × n: O(m × n) space
• A few variables for indices and temporary values: O(1) space

Total space complexity: O(m × n) + O(m × n) = O(m × n)

Where m and n are the number of rows and columns in the matrix, respectively.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Index Confusion Between Prefix Sum and Original Matrix

The Pitfall: One of the most common mistakes is mixing up the indexing between the prefix sum array (which has dimensions (m+1) x (n+1)) and the original matrix (which has dimensions m x n). The prefix sum array is intentionally padded with an extra row and column of zeros, so prefix_sum[i][j] corresponds to the sum up to mat[i-1][j-1].

Example of Incorrect Code:

# WRONG: Using same indices for both arrays
prefix_sum[i][j] = prefix_sum[i-1][j] + prefix_sum[i][j-1] - prefix_sum[i-1][j-1] + mat[i][j]

The Solution: Always remember that when building the prefix sum array, you need to offset the matrix indices:

# CORRECT: Offset mat indices by -1
prefix_sum[i][j] = prefix_sum[i-1][j] + prefix_sum[i][j-1] - prefix_sum[i-1][j-1] + mat[i-1][j-1]

2. Incorrect Block Boundary Calculation

The Pitfall: When calculating the block sum, forgetting to add 1 when using the prefix sum array indices can lead to off-by-one errors. Since prefix_sum[i][j] represents the sum up to position (i-1, j-1) in the original matrix, to include row r and column c, you need to use prefix_sum[r+1][c+1].

Example of Incorrect Code:

# WRONG: Forgetting to add 1 to boundaries
result[i][j] = (prefix_sum[bottom_row][right_col] - 
                prefix_sum[top_row][right_col] - 
                prefix_sum[bottom_row][left_col] + 
                prefix_sum[top_row][left_col])

The Solution: Always add 1 to the bottom and right boundaries when accessing the prefix sum array:

# CORRECT: Add 1 to bottom_row and right_col
result[i][j] = (prefix_sum[bottom_row + 1][right_col + 1] - 
                prefix_sum[top_row][right_col + 1] - 
                prefix_sum[bottom_row + 1][left_col] + 
                prefix_sum[top_row][left_col])

3. Boundary Validation Order

The Pitfall: Using min and max in the wrong order or with incorrect values when determining block boundaries.

Example of Incorrect Code:

# WRONG: Using matrix dimensions incorrectly
top_row = max(i - k, 0)
bottom_row = min(m, i + k)  # Should be m-1, not m

The Solution: Remember that matrix indices are 0-based, so the maximum valid index is m-1 for rows and n-1 for columns:

# CORRECT: Use m-1 and n-1 as maximum indices
top_row = max(i - k, 0)
left_col = max(j - k, 0)
bottom_row = min(i + k, rows - 1)
right_col = min(j + k, cols - 1)