Teaching Kids Programming – Matrix Prefix Sum Algorithm

Teaching Kids Programming: Videos on Data Structures and Algorithms

Given a two-dimensional integer matrix, return a new matrix A of the same dimensions where each element is set to A[i][j] = sum(matrix[r][c]) for all r ≤ i, c ≤ j.

Constraints
n, m ≤ 250 where n and m are the number of rows and columns in matrix
matrix[i][j] ≤ 2**12
Example 1
Input

matrix = [
    [2, 3],
    [5, 7]
]

Output

[
    [2, 5],
    [7, 17]
]

Algorithm to Compute the Prefix Sum for a Matrix

Matrix Prefix Sum is the 2D version of the accumulate function. For 1D version, which is the prefix sum for the array/lists, we can do the following:

def prefixSum(arr):
  for i in range(1, len(arr)):
    arr[i] += arr[i - 1]
  return ans

We can also use the accumulate from the itertools:

from itertools import accumulate
data = [1, 2, 3, 4]
# prefixSum = [1, 3, 6, 10]
prefixSum = list(accumulate(data))

For the 2D Matrix, we can do this for each rows, and then for each columns:

class Solution:
    def matrixPrefixSum(self, matrix):
        if not matrix:
            return []

        row = len(matrix)
        if 0 == row:
            return matrix
        col = len(matrix[0])
        if 0 == col:
            return matrix

        for r in range(row):
            matrix[r] = list(accumulate(matrix[r]))
            # or use the following            
            # for c in range(1, col):
            #     matrix[r][c] += matrix[r][c - 1]

        for r in range(1, row):
            for c in range(col):            
                matrix[r][c] += matrix[r - 1][c]

        return matrix        

The time complexity is O(RC) where R and C are the number of rows and columns for the Matrix. The space complexity is O(1) constant as we are building the prefix sum matrix in place.

–EOF (The Ultimate Computing & Technology Blog) —