1314. Matrix Block Sum


Problem Description

Given a matrix mat with dimensions m x n and an integer k, the task is to compute a new matrix answer of the same dimensions. For each cell answer[i][j] in the new matrix, we must calculate the sum of all elements within a square block with side length 2k + 1, centered at mat[i][j]. The boundaries of the square are constrained by the limits of the array, meaning we cannot go outside the matrix when calculating the sum. The result for answer[i][j] should only include the sum of valid elements from mat that lie within the specified range and within the matrix.

Intuition

The brute-force approach to solving this problem would involve iterating over each cell of the matrix, and for each cell, iteratively summing all the elements in the K-radius square. This would require a large number of redundant computations and would not be efficient, particularly for large matrices or values of K.

To arrive at a more efficient solution, we can use dynamic programming, specifically the concept of 2D prefix sums. The idea here is to preprocess the input matrix to generate a prefix sum matrix where each element pre[i][j] stores the sum of all elements in the rectangle defined by the upper left corner (0,0) and the lower right corner (i-1, j-1) of the original matrix. With this prefix sum matrix, we can compute the sum of any submatrix in constant time, using the inclusion-exclusion principle.

The steps to implement this are as follows:

  1. Build the prefix sum matrix (pre) where the sum up to each point (i, j) is calculated by adding the current element to the sum of the elements above and to the left, subtracting the sum of the overlapping corner (already counted twice).

  2. Define a helper function get(i, j) to safely retrieve the value of pre[i][j], handling boundary cases to stay within the matrix dimensions.

  3. Loop through each cell (i, j) in the original matrix, and for each cell, calculate the sum of the K-radius square around it using the prefix sum matrix and the helper function. This involves adding and subtracting values from the four corners of the block surrounding (i, j), constrained by k.

The result is an efficient algorithm that computes the desired answer matrix with much fewer operations compared to the brute-force method, thus making it scalable to larger matrices and values of k. The overall time complexity is O(m * n) for the sum matrix construction plus O(m * n) for calculating the block sums, and the space complexity is O(m * n) for storing the prefix sums.

Learn more about Prefix Sum patterns.

Solution Approach

The solution to the problem uses dynamic programming with a 2D prefix sum matrix to optimize the process of computing the sums of submatrices. Let's break down the steps taken in the provided solution code:

  1. Initial Prefix Sum Matrix Generation: A 2D prefix sum matrix pre is generated where pre[i][j] will contain the sum of all elements from the top-left corner (0, 0) to the position (i-1, j-1) in the original matrix mat. This matrix pre is of size (m + 1) x (n + 1) to handle the prefix sum calculation conveniently, even for the cells on the topmost row and leftmost column.

    The prefix sum is calculated as follows:

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

    This equation accounts for the current element and the sums above and to the left, removing the overlapped corner that was added twice.

  2. Safe Prefix Sum Retrieval Function (get): The get function is defined to safely retrieve values from the pre matrix, taking care of the boundaries, and ensuring that indices outside the matrix bounds map to the nearest valid index. This function ensures that when we perform the sum calculation, we do not go outside the matrix.

    1def get(i, j):
    2    i = max(min(m, i), 0)
    3    j = max(min(n, j), 0)
    4    return pre[i][j]
  3. Computing the Block Sums: We then iterate over each cell (i, j) in the matrix mat. The sum for each answer[i][j] is computed using the get function to obtain values from the prefix sum matrix corresponding to the corners of the K-radius block:

    1ans[i][j] = (
    2    get(i + k + 1, j + k + 1)
    3    - get(i + k + 1, j - k)
    4    - get(i - k, j + k + 1)
    5    + get(i - k, j - k)
    6)

    This step applies the inclusion-exclusion principle to subtract the sums outside the K-radius block by considering the bottom right corner (i + k + 1, j + k + 1), the top right (i + k + 1, j - k), top left (i - k, j - k), and bottom left corners (i - k, j + k + 1).

The matrixBlockSum function ultimately returns the computed ans matrix, which represents the sum of elements within k distance from each element in the original mat.

The code structure provided makes use of a 2D prefix sum matrix and carefully handles edge cases with the get function, ensuring the solution is clean, efficient, and robust for the given problem.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Problem: Given a list of tasks and a list of requirements, compute a sequence of tasks that can be performed, such that we complete every task once while satisfying all the requirements.

Which of the following method should we use to solve this problem?

Example Walkthrough

Let's use a simple 3x3 matrix as an example to illustrate how the solution approach works, with k=1:

Original matrix mat:

1| 1 | 2 | 3 |
2| 4 | 5 | 6 |
3| 7 | 8 | 9 |
  1. Initial Prefix Sum Matrix Generation: We first build an enlarged prefix sum matrix pre of size 4x4 to include a buffer for handling the edges. The construction would look like this after using the described formula:

    Initial pre matrix filled with zeros (for illustration):

    1| 0 | 0 | 0 | 0 |
    2| 0 |   |   |   |
    3| 0 |   |   |   |
    4| 0 |   |   |   |

    After calculating the prefix sums:

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

    You can see that pre[3][3] has a value of 45, which is the sum of all elements of mat.

  2. Safe Prefix Sum Retrieval Function (get): The get function is used to safely access elements in the pre matrix, ensuring that we do not access indices outside the bounds of the matrix.

  3. Computing the Block Sums: We compute the sums for each cell in ans considering k=1. For the central element ans[1][1] corresponding to mat[1][1] which is 5, the calculation using the corners would be:

    1ans[1][1] = (
    2    get(1 + 1 + 1, 1 + 1 + 1)
    3    - get(1 + 1 + 1, 1 - 1)
    4    - get(1 - 1, 1 + 1 + 1)
    5    + get(1 - 1, 1 - 1)
    6)
    7ans[1][1] = pre[3][3] - pre[3][1] - pre[1][3] + pre[1][1]
    8ans[1][1] = 45 - 12 - 6 + 1
    9ans[1][1] = 28

    The get function ensures we don't index into pre at negative indices or beyond its size by clamping the values appropriately.

    The complete ans matrix would be computed similarly, for example, the top-left corner ans[0][0] would be computed as:

    1ans[0][0] = (
    2    get(0 + 1 + 1, 0 + 1 + 1)
    3    - get(0 + 1 + 1, 0 - 1)
    4    - get(0 - 1, 0 + 1 + 1)
    5    + get(0 - 1, 0 - 1)
    6)
    7ans[0][0] = pre[2][2] - pre[2][0] - pre[0][2] + pre[0][0]
    8ans[0][0] = 12 - 0 - 0 + 0
    9ans[0][0] = 12

    Which is the sum of the 2x2 block that is within k distance from the top-left element 1.

Final matrix ans after computing sums for all other cells:

1| 12 | 21 | 16 |
2| 27 | 45 | 33 |
3| 24 | 39 | 28 |

This walk-through demonstrates the computational steps of the algorithm using the example matrix. Notice how the sums reflect the areas around each cell, constrained by k and the bounds of the matrix.

Solution Implementation

1from typing import List
2
3class Solution:
4    def matrixBlockSum(self, mat: List[List[int]], k: int) -> List[List[int]]:
5        # Get the dimensions of the matrix
6        num_rows, num_cols = len(mat), len(mat[0])
7      
8        # Initialize a prefix sum matrix with one extra row and column
9        prefix_sum = [[0] * (num_cols + 1) for _ in range(num_rows + 1)]
10      
11        # Populate the prefix sum matrix
12        for row in range(1, num_rows + 1):
13            for col in range(1, num_cols + 1):
14                prefix_sum[row][col] = (
15                    prefix_sum[row - 1][col] +
16                    prefix_sum[row][col - 1] -
17                    prefix_sum[row - 1][col - 1] +
18                    mat[row - 1][col - 1]
19                )
20
21        # Helper function to get the cumulative sum up to (i, j) in the prefix sum matrix
22        def get_cumulative_sum(i, j):
23            # Boundary check to ensure indices fall within the valid range
24            i = max(min(num_rows, i), 0)
25            j = max(min(num_cols, j), 0)
26            return prefix_sum[i][j]
27
28        # Initialize the answer matrix with zeros
29        answer = [[0] * num_cols for _ in range(num_rows)]
30      
31        # Compute the block sum for each element in the matrix
32        for row in range(num_rows):
33            for col in range(num_cols):
34                answer[row][col] = (
35                    get_cumulative_sum(row + k + 1, col + k + 1) -
36                    get_cumulative_sum(row + k + 1, col - k) -
37                    get_cumulative_sum(row - k, col + k + 1) +
38                    get_cumulative_sum(row - k, col - k)
39                )
40              
41        # Return the answer matrix
42        return answer
43
1class Solution {
2    private int[][] prefixSum; // 2D array to hold the prefix sum of matrix
3    private int numRows;
4    private int numCols;
5
6    /**
7     * Calculates the block sum of each element in the matrix surrounded by 'k' distance.
8     * @param mat The original matrix of integers
9     * @param k The given distance around the element to calculate sum
10     * @return The matrix representing the block sum of each element
11     */
12    public int[][] matrixBlockSum(int[][] mat, int k) {
13        numRows = mat.length; // number of rows in the input matrix
14        numCols = mat[0].length; // number of columns in the input matrix
15      
16        // Define a new matrix with additional row and column to easily compute prefix sum
17        prefixSum = new int[numRows + 1][numCols + 1];
18      
19        // Populating the prefix sum matrix
20        for (int i = 1; i <= numRows; ++i) {
21            for (int j = 1; j <= numCols; ++j) {
22                // Each cell is the sum of values above it, to its left, 
23                // minus the cell diagonally up and left to remove the double count, plus the current cell
24                prefixSum[i][j] = prefixSum[i - 1][j] + prefixSum[i][j - 1] - prefixSum[i - 1][j - 1] + mat[i - 1][j - 1];
25            }
26        }
27      
28        // Create a matrix to hold answers
29        int[][] answer = new int[numRows][numCols];
30
31        // Calculating the block sum for each cell
32        for (int i = 0; i < numRows; ++i) {
33            for (int j = 0; j < numCols; ++j) {
34                // Calculate the sum of the block using prefix sum
35                answer[i][j] = getSum(i + k + 1, j + k + 1) - getSum(i + k + 1, j - k) - getSum(i - k, j + k + 1) + getSum(i - k, j - k);
36            }
37        }
38        return answer;
39    }
40
41    /**
42     * Get the prefix sum of the cell (i, j) ensuring the indices are within bounds.
43     * @param i The row index in the prefix sum matrix
44     * @param j The column index in the prefix sum matrix
45     * @return The prefix sum value of the cell (i, j) bounded properly
46     */
47    private int getSum(int i, int j) {
48        // Bound the indices within the limits of the prefixSum matrix
49        i = Math.max(Math.min(numRows, i), 0);
50        j = Math.max(Math.min(numCols, j), 0);
51        return prefixSum[i][j];
52    }
53}
54
1class Solution {
2public:
3    vector<vector<int>> matrixBlockSum(vector<vector<int>>& mat, int k) {
4        int rows = mat.size(), cols = mat[0].size();
5        vector<vector<int>> prefixSum(rows + 1, vector<int>(cols + 1));
6
7        // Compute 2D prefix sum
8        for (int i = 1; i <= rows; ++i) {
9            for (int j = 1; j <= cols; ++j) {
10                prefixSum[i][j] = prefixSum[i - 1][j] 
11                                + prefixSum[i][j - 1] 
12                                - prefixSum[i - 1][j - 1] 
13                                + mat[i - 1][j - 1];
14            }
15        }
16
17        vector<vector<int>> answer(rows, vector<int>(cols));
18
19        // Iterate through all the cells of the matrix
20        for (int i = 0; i < rows; ++i) {
21            for (int j = 0; j < cols; ++j) {
22                // Calculate block sum using prefixSum
23                answer[i][j] = getBlockSum(i + k + 1, j + k + 1, rows, cols, prefixSum) 
24                             - getBlockSum(i + k + 1, j - k, rows, cols, prefixSum) 
25                             - getBlockSum(i - k, j + k + 1, rows, cols, prefixSum) 
26                             + getBlockSum(i - k, j - k, rows, cols, prefixSum);
27            }
28        }
29
30        return answer;
31    }
32
33    int getBlockSum(int row, int col, int maxRow, int maxCol, vector<vector<int>>& prefixSum) {
34        // Clamp row and column indices to the valid range
35        row = max(min(maxRow, row), 0);
36        col = max(min(maxCol, col), 0);
37      
38        // Return the value at the bounded prefix sum location
39        return prefixSum[row][col];
40    }
41};
42
1// Define the matrix type as an array of arrays of numbers
2type Matrix = number[][];
3
4// Calculate the 2D prefix sum matrix
5function buildPrefixSumMatrix(mat: Matrix): Matrix {
6    const rows = mat.length, cols = mat[0].length;
7    const prefixSum: Matrix = Array.from({ length: rows + 1 }, () => Array(cols + 1).fill(0));
8
9    // Compute the 2D prefix sum
10    for (let i = 1; i <= rows; ++i) {
11        for (let j = 1; j <= cols; ++j) {
12            prefixSum[i][j] = prefixSum[i - 1][j] + prefixSum[i][j - 1] 
13                            - prefixSum[i - 1][j - 1] + mat[i - 1][j - 1];
14        }
15    }
16  
17    return prefixSum;
18}
19
20// Given the row and column indices, return the prefix sum within the bounded indices
21function getBlockSum(row: number, col: number, maxRow: number, maxCol: number, prefixSum: Matrix): number {
22    // Clamp row and column indices to the valid range
23    row = Math.max(Math.min(maxRow, row), 0);
24    col = Math.max(Math.min(maxCol, col), 0);
25  
26    // Return the value at the bounded prefix sum location
27    return prefixSum[row][col];
28}
29
30// Computes the matrix block sum given the original matrix and the value k
31function matrixBlockSum(mat: Matrix, k: number): Matrix {
32    const rows = mat.length, cols = mat[0].length;
33    const prefixSum = buildPrefixSumMatrix(mat);
34    const answer: Matrix = Array.from({ length: rows }, () => Array(cols).fill(0));
35
36    // Iterate through all the cells of the matrix
37    for (let i = 0; i < rows; ++i) {
38        for (let j = 0; j < cols; ++j) {
39            // Calculate block sum using prefixSum
40            answer[i][j] = getBlockSum(i + k + 1, j + k + 1, rows, cols, prefixSum)
41                         - getBlockSum(i + k + 1, j - k, rows, cols, prefixSum)
42                         - getBlockSum(i - k, j + k + 1, rows, cols, prefixSum)
43                         + getBlockSum(i - k, j - k, rows, cols, prefixSum);
44        }
45    }
46
47    return answer;
48}
49

Time and Space Complexity

Time Complexity:

The time complexity of the given code can be analyzed in two major parts: the construction of the prefix sum matrix pre and the computation of the answer matrix ans.

  1. Prefix Sum Matrix Computation: Constructing the prefix sum matrix requires iterating through each element of the original m x n matrix mat, and for each element, a constant amount of work is performed (adding and subtracting values). Therefore, the time complexity for this part is O(m * n).

  2. Answer Matrix Computation: The computation of the answer matrix also involves iterating through all elements in the m x n matrix and performing a constant amount of work for each to calculate the sum of the block using the prefix sum matrix. Hence, the time complexity for this part is also O(m * n).

Combining both parts, the overall time complexity remains O(m * n) because both parts are sequentially executed and have the same order of complexity.

Space Complexity:

The space complexity can be analyzed as follows:

  1. Prefix Sum Matrix: An additional m+1 x n+1 matrix pre is used for the prefix sums, which gives us a space complexity of O((m + 1) * (n + 1)).

  2. Answer Matrix: The answer matrix ans is of the same size as the input matrix mat, i.e., m x n. Since it is used to store the final output, it might not be counted towards additional space in some analysis contexts. But for a comprehensive analysis, we consider it and, therefore, O(m * n).

When combined, the term O((m + 1) * (n + 1)) is dominant. However, since the +1 is a constant and does not change the order of growth, the overall space complexity simplifies to O(m * n).

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


Fast Track Your Learning with Our Quick Skills Quiz:

What is the worst case running time for finding an element in a binary search tree(not necessarily balanced) of size n?


Recommended Readings


Got a question? Ask the Monster Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.

Tired of the LeetCode Grind?

Our structured approach teaches you the patterns behind problems, so you can confidently solve any challenge. Get started now to land your dream tech job.

Get Started

🪄