LeetCode 1314 - Matrix Block Sum

Difficulty: 🟡 Medium
Topics: Array, Matrix, Prefix Sum

Solution

Problem Understanding

The problem requires computing a matrix block sum. Given a matrix mat of size m x n and an integer k, the task is to produce a new matrix answer of the same dimensions. Each element answer[i][j] is defined as the sum of all elements mat[r][c] that lie within a square block centered at (i, j) with radius k. The block is defined by the conditions i - k <= r <= i + k and j - k <= c <= j + k, but the block must remain within valid matrix indices.

The input matrix is guaranteed to have dimensions 1 <= m, n <= 100 and values between 1 and 100. The integer k satisfies 1 <= k <= 100, meaning the block can potentially extend beyond the matrix bounds, so we must carefully handle boundary conditions to avoid index errors.

Key edge cases include blocks at the corners or edges of the matrix, where part of the block lies outside the matrix. A naive approach that iterates over every element in the block for every matrix position can become inefficient for large matrices because of repeated summations.

The goal is an efficient solution that calculates these sums without excessive recomputation.

Approaches

Brute Force Approach

The brute-force solution iterates over each element (i, j) in the matrix and sums all elements in its k-radius block. For each answer[i][j], we check all valid r and c indices within i-k to i+k and j-k to j+k. While correct, this approach is inefficient. For a matrix of size m x n and a block of size up to (2k+1)^2, the time complexity is O(m * n * (2k+1)^2). Given m, n, k <= 100, this could require up to 100 million operations, which is borderline slow for online coding platforms.

Optimal Approach: Prefix Sum

The key insight is that we can use a 2D prefix sum matrix to compute any submatrix sum in constant time. A prefix sum matrix psum is defined such that psum[i][j] represents the sum of all elements from (0,0) to (i-1,j-1) (1-based offset to simplify boundaries). With this matrix, the sum of any rectangle (r1,c1) to (r2,c2) can be computed as:

sum = psum[r2+1][c2+1] - psum[r1][c2+1] - psum[r2+1][c1] + psum[r1][c1]

This approach allows each answer[i][j] to be calculated in O(1) time after building the prefix sum, reducing the overall time complexity to O(m * n).

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(m * n * k^2)	O(1)	Sum each block individually
Optimal	O(m * n)	O(m * n)	Use 2D prefix sum for fast rectangle sum

Algorithm Walkthrough

Build the Prefix Sum Matrix: Create a matrix psum of size (m+1) x (n+1) initialized to zero. For each element (i,j) in mat, compute psum[i+1][j+1] = psum[i][j+1] + psum[i+1][j] - psum[i][j] + mat[i][j]. This ensures psum[i+1][j+1] contains the sum of all elements from (0,0) to (i,j).
Compute Answer Matrix: Initialize answer as a zero matrix of size m x n. For each element (i,j), define the boundaries of the block:

r1 = max(0, i-k)
c1 = max(0, j-k)
r2 = min(m-1, i+k)
c2 = min(n-1, j+k)

Use the prefix sum to calculate the sum of this block:

answer[i][j] = psum[r2+1][c2+1] - psum[r1][c2+1] - psum[r2+1][c1] + psum[r1][c1]

Return the Answer: After filling all elements, return the answer matrix.

Why it works: The prefix sum ensures that every rectangular block sum can be computed in constant time using inclusion-exclusion. Boundary conditions are correctly handled by clamping indices, so no invalid access occurs.

Python Solution

from typing import List

class Solution:
    def matrixBlockSum(self, mat: List[List[int]], k: int) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        psum = [[0] * (n + 1) for _ in range(m + 1)]

        # Build prefix sum
        for i in range(m):
            for j in range(n):
                psum[i+1][j+1] = psum[i][j+1] + psum[i+1][j] - psum[i][j] + mat[i][j]

        # Compute answer using prefix sum
        answer = [[0] * n for _ in range(m)]
        for i in range(m):
            for j in range(n):
                r1, c1 = max(0, i - k), max(0, j - k)
                r2, c2 = min(m - 1, i + k), min(n - 1, j + k)
                answer[i][j] = psum[r2+1][c2+1] - psum[r1][c2+1] - psum[r2+1][c1] + psum[r1][c1]

        return answer

The code first constructs a prefix sum matrix to allow O(1) sum queries for any rectangle. Then, for each matrix element, it calculates the valid block boundaries and applies the prefix sum formula. Clamping with max and min handles edges and corners.

Go Solution

func matrixBlockSum(mat [][]int, k int) [][]int {
    m, n := len(mat), len(mat[0])
    psum := make([][]int, m+1)
    for i := range psum {
        psum[i] = make([]int, n+1)
    }

    // Build prefix sum
    for i := 0; i < m; i++ {
        for j := 0; j < n; j++ {
            psum[i+1][j+1] = psum[i][j+1] + psum[i+1][j] - psum[i][j] + mat[i][j]
        }
    }

    // Compute answer using prefix sum
    answer := make([][]int, m)
    for i := range answer {
        answer[i] = make([]int, n)
        for j := range answer[i] {
            r1, c1 := max(0, i-k), max(0, j-k)
            r2, c2 := min(m-1, i+k), min(n-1, j+k)
            answer[i][j] = psum[r2+1][c2+1] - psum[r1][c2+1] - psum[r2+1][c1] + psum[r1][c1]
        }
    }

    return answer
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

The Go implementation mirrors Python closely. The key differences are explicit slice allocation, use of helper max and min functions, and zero-based indexing handled via m+1 and n+1 in the prefix sum matrix.

Worked Examples

Example 1:

mat = [[1,2,3],[4,5,6],[7,8,9]], k = 1

Compute prefix sum:

	1	2	3
0	0	0	0
1	1	3	6
2	5	12	21
3	12	27	45

Calculate answer[0][0]: r1=0,c1=0,r2=1,c2=1 → sum = psum[2][2] - psum[0][2] - psum[2][0] + psum[0][0] = 12.

Compute all elements similarly to get:

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

Example 2:

mat = [[1,2,3],[4,5,6],[7,8,9]], k = 2

Block extends beyond matrix, sum of all elements = 45 for every