1277. Count Square Submatrices with All Ones

Given a m * n matrix of ones and zeros, return how many square submatrices have all ones.

Example 1:

Input: matrix =
[
 [0,1,1,1],
 [1,1,1,1],
 [0,1,1,1]
]
Output: 15
Explanation:
There are 10 squares of side 1.
There are 4 squares of side 2.
There is  1 square of side 3.
Total number of squares = 10 + 4 + 1 = 15.

Example 2:

Input: matrix = 
[
  [1,0,1],
  [1,1,0],
  [1,1,0]
]
Output: 7
Explanation:
There are 6 squares of side 1.  
There is 1 square of side 2. 
Total number of squares = 6 + 1 = 7.

Constraints:

• 1 <= arr.length <= 300
• 1 <= arr[0].length <= 300
• 0 <= arr[i][j] <= 1

Approach 01:

class Solution {
public:
    int countSquares(vector<vector<int>>& matrix) {
        int rows = matrix.size();
        int cols = matrix[0].size();

        // Update the matrix with the size of square that can be formed at each
        // cell
        for (int row = 0; row < rows; ++row) {
            for (int col = 0; col < cols; ++col) {
                if (matrix[row][col] == 1 && row > 0 && col > 0)
                    matrix[row][col] += min({matrix[row - 1][col - 1], matrix[row - 1][col],  matrix[row][col - 1]});
            }
        }

        // Calculate the total number of squares
        return accumulate(matrix.begin(), matrix.end(), 0, [](int totalSquares, const vector<int>& currentRow) {
            return totalSquares + accumulate(currentRow.begin(), currentRow.end(), 0); 
        });
    }
};

class Solution:
    def countSquares(self, matrix: list[list[int]]) -> int:
        rows = len(matrix)
        cols = len(matrix[0])

        # Update the matrix with the size of square that can be formed at each cell
        for row in range(rows):
            for col in range(cols):
                if matrix[row][col] == 1 and row > 0 and col > 0:
                    matrix[row][col] += min(matrix[row - 1][col - 1], matrix[row - 1][col], matrix[row][col - 1])

        # Calculate the total number of squares
        return sum(sum(currentRow) for currentRow in matrix)

Time Complexity

• Iterating through the matrix:

  The matrix is processed cell by cell in a nested loop, with one loop iterating over the rows and another over the columns. Since there are \(r\) rows and \(c\) columns, this takes \(O(r \times c)\) time.

• Accumulating the results:

  To calculate the total number of squares, two accumulate calls are made. The first one sums up each row, which takes \(O(c)\) for each of the \(r\) rows, and the second sums up the total number of squares from these row sums. So, this takes \(O(r \times c)\) as well.

• Overall Time Complexity:

  Both the matrix processing and accumulation contribute \(O(r \times c)\), so the overall time complexity is \(O(r \times c)\).

Space Complexity

• In-place updates to the matrix:

  The algorithm modifies the input matrix directly, so no additional 2D space is used. This gives a space complexity of \(O(1)\) for extra space, except for the input itself.

• Accumulation variables:

  The space used by temporary variables like totalSquares and the parameters for accumulate is constant.

• Overall Space Complexity:

  Since the algorithm uses constant space apart from the input matrix, the space complexity is \(O(1)\).