Matrix Block Sum

Given a m * n matrix mat and an integer K, return a matrix answer where each answer[i][j] is the sum of all elements mat[r][c] for i - K <= r <= i + K, j - K <= c <= j + K, and (r, c) is a valid position in 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]]

Approach

Java

import java.util.Arrays;

public class MatrixBlockSum {
    public static void main(String[] args) {
        int mat[][] = { { 1, 2, 3 }, { 4, 5, 6 }, { 7, 8, 9 } };
        int K = 1;
        mat = matrixBlockSum(mat, K);
        for (int i = 0; i < mat.length; i++) {
            System.out.println(Arrays.toString(mat[i]));
        }
        
    }

    // method to find the block sum of the matrix
    static int[][] matrixBlockSum(int[][] mat, int K) {
        int n = mat.length;
        int m = mat[0].length;

        int[][] res = mat;
        int[][] ans = new int[n][m];
        for (int i = 0; i < n; i++) {
            for (int j = 1; j < m; j++)
                res[i][j] += res[i][j - 1];
        }
        for (int i = 0; i < m; i++)
            for (int j = 1; j < n; j++)
                res[j][i] += res[j - 1][i];
        
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) {
                int col = Math.min(m - 1, j + K);
                int row = Math.min(n - 1, i + K);
                ans[i][j] = res[row][col];
                if (i - K - 1 >= 0)
                    ans[i][j] -= res[i - K - 1][col];
                if (j - K - 1 >= 0)
                    ans[i][j] -= res[row][j - K - 1];
                if (i - K - 1 >= 0 && j - K - 1 >= 0)
                    ans[i][j] += res[i - K - 1][j - K - 1];
            }
        }
        return ans;
    }
}

C++

#include <bits/stdc++.h>
using namespace std;

//function to find the block sum of the matrix
vector<vector<int>> matrixBlockSum(vector<vector<int>>& mat, 
int K) 
{
        int n=mat.size();
        int m=mat[0].size();
        
        vector<vector<int>> res=mat; 
        vector<vector<int>> ans(n,vector<int>(m,0));
        for(int i=0;i<n;i++)
        {
            for(int j=1;j<m;j++)
                  res[i][j]+=res[i][j-1];
        }
        for(int i=0;i<m;i++)
              for(int j=1;j<n;j++)
                     res[j][i]+=res[j-1][i];
        for(int i=0;i<n;i++)
        {
            for(int j=0;j<m;j++)
            {
                int col=min(m-1,j+K),row=min(n-1,i+K);
                ans[i][j]=res[row][col];
                if(i-K-1>=0)
                      ans[i][j]-=res[i-K-1][col];
               if(j-K-1>=0)
                      ans[i][j]-=res[row][j-K-1];
               if(i-K-1>=0&&j-K-1>=0)
                      ans[i][j]+=res[i-K-1][j-K-1];
            }
        }
        return ans;
}
int main()
{
  vector<vector<int>> mat ={{1,2,3},
                            {4,5,6},
                            {7,8,9}};
  int K = 1;
  mat=matrixBlockSum(mat,K);    
  cout<<"[";
  for(int i=0;i<mat.size();i++)
    {
       cout<<"[";
       for(int j=0;j<mat[i].size();j++)
         {
             if(j!=mat[i].size()-1)
               cout<<mat[i][j]<<",";
             else
              cout<<mat[i][j];
         }
         if(i!=mat.size()-1)
            cout<<"],";
         else
          cout<<"]";
    }
  cout<<"]";
  return 0;
}