Matrix Block Sum — Medium Problem & Solution

Return a matrix answer of the same size as mat where answer[i][j] is the sum of all mat[r][c] with i - k <= r <= i + k and j - k <= c <= j + k, counting…

  • Difficulty: Medium
  • Topics: Arrays, Matrix, Prefix Sum
  • Asked at: Amazon, Google, Adobe
  • Time limit: 2 s
  • Memory limit: 256 MB
  • Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

Return a matrix answer of the same size as mat where answer[i][j] is the sum of all mat[r][c] with i - k <= r <= i + k and j - k <= c <= j + k, counting only positions inside 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]]
Explanation: A radius of 2 covers the whole matrix from every cell.

Example 3

Input: mat = [[5]], k = 3
Output: [[5]]

Constraints

  • m == mat.length
  • n == mat[i].length
  • 1 <= m, n, k <= 100
  • 1 <= mat[i][j] <= 100

How to solve Matrix Block Sum

Precompute a 2D prefix-sum table, then answer every block with the standard inclusion-exclusion of four corners. Clamping handles blocks that hang over the edges.

Approach

  1. Build pre with an extra zero row and column, where pre[i+1][j+1] is the sum of mat[0…i][0…j].
  2. For each cell, clamp the block to [max(0, i-k), min(m-1, i+k)] and likewise for columns.
  3. The block sum is pre[r2+1][c2+1] - pre[r1][c2+1] - pre[r2+1][c1] + pre[r1][c1].

Why it works

The four-corner formula subtracts the two overlapping strips and adds back the doubly-subtracted corner — plain inclusion-exclusion. The padding row and column remove every boundary special case, which is what makes the clamped indices safe to use directly.

Complexity

  • Time — O(m · n)
  • Space — O(m · n)

Pitfalls

  • Forgetting the + pre[r1][c1] term double-subtracts the overlap.
  • The clamp is on the block, not on the prefix indices — the padded table then handles the rest.
  • The totals reach 100 · 100 · 100 = 10^6, comfortably inside int.

Reference solution

Python

from typing import List

def matrixBlockSum(mat: List[List[int]], k: int) -> List[List[int]]:
    m, n = len(mat), len(mat[0])
    pre = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(m):
        for j in range(n):
            pre[i + 1][j + 1] = mat[i][j] + pre[i][j + 1] + pre[i + 1][j] - pre[i][j]
    out = []
    for i in range(m):
        row = []
        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)
            row.append(pre[r2 + 1][c2 + 1] - pre[r1][c2 + 1] - pre[r2 + 1][c1] + pre[r1][c1])
        out.append(row)
    return out

JavaScript

var matrixBlockSum = function(mat, k) {
    var m = mat.length, n = mat[0].length, i, j;
    var pre = [];
    for (i = 0; i <= m; i++) {
        var prow = [];
        for (j = 0; j <= n; j++) prow.push(0);
        pre.push(prow);
    }
    for (i = 0; i < m; i++) {
        for (j = 0; j < n; j++) {
            pre[i + 1][j + 1] = mat[i][j] + pre[i][j + 1] + pre[i + 1][j] - pre[i][j];
        }
    }
    var out = [];
    for (i = 0; i < m; i++) {
        var row = [];
        for (j = 0; j < n; j++) {
            var r1 = Math.max(0, i - k), c1 = Math.max(0, j - k);
            var r2 = Math.min(m - 1, i + k), c2 = Math.min(n - 1, j + k);
            row.push(pre[r2 + 1][c2 + 1] - pre[r1][c2 + 1] - pre[r2 + 1][c1] + pre[r1][c1]);
        }
        out.push(row);
    }
    return out;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 667 arrays problems · the whole catalogue