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: 1515
Explanation:
There are 1010 squares of side 11.
There are 44 squares of side 22.
There is 11 square of side 33.
Total number of squares =10+4+1=15= 10 + 4 + 1 = 15.

Example 2:

Input:

matrix =
[
  [1,0,1],
  [1,1,0],
  [1,1,0]
]

Output: 77
Explanation:
There are 66 squares of side 11.
There is 11 square of side 22.
Total number of squares =6+1=7= 6 + 1 = 7.

Constraints:

  • 1ā‰¤1 \leq arr.length ā‰¤300\leq 300
  • 1ā‰¤1 \leq arr[0].length ā‰¤300\leq 300
  • 0ā‰¤0 \leq arr[i][j] ā‰¤1\leq 1

Solution

Brute Force

First, let's say that a square is located at a cell (i,j) if its bottom right corner is located at that cell.

We can observe that if there is a length xx square with all ones located at (i,j), then there exists squares with all ones that have lengths xāˆ’1,xāˆ’2ā€¦,3,2,1x-1,x-2 \dots, 3,2,1. This is because if a square submatrix with length xx located there exists, then square submatrices with lengths xāˆ’1,xāˆ’2ā€¦3,2,1x-1, x-2 \dots 3,2,1 must exist as well.

For each cell (i,j) we'll find the length of the largest square submatrix that has all ones with its bottom right corner located at (i,j). Let's denote this value as xi,jx_{i,j}. We can observe that the number of square submatrices with all ones and its bottom right corner located at (i,j) is simply xi,jx_{i,j}. To find xi,jx_{i,j}, we can brute force through all possible lengths and check if a square submatrix with that respective length exists. Once we find xi,jx_{i,j} for each cell (i,j), our final answer is the sum of all xi,jx_{i,j}.

Full Solution

Our full solution will involve dynamic programming.

Let dp[i][j] represent xi,jx_{i,j}.

Since dynamic programming uses the answers to sub-problems to calculate answers to a larger problem, let's try to see how we can use other values from dp to calculate dp[i][j] for some cell (i,j).

Let's say that a square submatrix is located at a cell (i,j) if its bottom right corner is located there.

First, let's assume dp[i][j] = k for some positive integer kk. This means that the largest square submatrix with all ones located at (i,j) has length kk. We can observe that this means there exists square submatrices with all ones at (i,j-1), (i-1,j), and (i-1,j-1) that have length kāˆ’1k-1. In addition, the cell (i,j) is also 11.

If we look at all the cells that are covered by the square submatrices at (i,j-1), (i-1,j), and (i-1,j-1) with length kāˆ’1k-1 and the cell at (i,j), we obtain a square with length kk located at (i,j).

Example

Here is a diagram with k=5k = 5 and the cell at (5,5) to help visualize this observation.

From the observation mentioned above, we can deduce that dp[i-1][j], dp[i][j-1], dp[i-1][j-1] ā‰„kāˆ’1\geq k-1.

For dp[i][j] to be greater or equal to kk, then (i,j) must be 11 and dp[i-1][j], dp[i][j-1], dp[i-1][j-1] all have to be greater or equal to kāˆ’1k-1.

This translates into dp[i][j] = min(dp[i-1][j],dp[i][j-1],dp[i-1][j-1]) + 1 if the cell (i,j) is a 11. For most cases, this is how we can calculate dp[i][j].

If the cell is in the first row or column however (i.e. i=0i = 0 or j=0j = 0), then dp[i][j] is the same value as the cell (i,j).

Obviously, if the cell (i,j) is 00, then dp[i][j] = 0.

Our final answer is just the sum of all values stored in dp.

Time Complexity

We can calculate the value of a cell in dp in O(1)\mathcal{O}(1) and since there are O(MN)\mathcal{O}(MN) cells, our time complexity is O(MN)\mathcal{O}(MN).

Time Complexity: O(MN)\mathcal{O}(MN)

Space Complexity

We store O(MN)\mathcal{O}(MN) cells in dp so our space complexity is O(MN)\mathcal{O}(MN).

Space Complexity: O(MN)\mathcal{O}(MN)

Bonus: We can use the space optimization mentioned in this article to optimize memory to O(N)O(N). Since we only use rows iāˆ’1i-1 and ii in dp to calculate dp values for row ii, we only need to maintain two rows of memory for dp, which is O(N)O(N).

C++ Solution

class Solution {
   public:
    int countSquares(vector<vector<int>>& matrix) {
        int m = matrix.size();
        int n = matrix[0].size();  // dimensions for matrix
        int ans = 0;
        vector<vector<int>> dp(m, vector<int>(n));
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (i == 0 || j == 0) {  // cell is in first row or column
                    dp[i][j] = matrix[i][j];
                } else if (matrix[i][j] == 1) {
                    dp[i][j] = min({dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1]}) + 1;
                }
                ans += dp[i][j];
            }
        }
        return ans;
    }
};

Java Solution

class Solution {
    public int countSquares(int[][] matrix) {
        int m = matrix.length;
        int n = matrix[0].length; // dimensions for matrix
        int[][] dp = new int[m][n];
        int ans = 0;
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (i == 0 || j == 0) { // cell is in first row or column
                    dp[i][j] = matrix[i][j];
                } else if (matrix[i][j] == 1) {
                    dp[i][j] = Math.min(dp[i - 1][j], Math.min(dp[i][j - 1], dp[i - 1][j - 1])) + 1;
                }
                ans += dp[i][j];
            }
        }
        return ans;
    }
}

Python Solution

class Solution:
    def countSquares(self, matrix: List[List[int]]) -> int:
        m = len(matrix)
        n = len(matrix[0])  # dimensions for matrix
        dp = [[0] * n for a in range(m)]
        ans = 0
        for i in range(m):
            for j in range(n):
                if i == 0 or j == 0:  # cell is in first row or column
                    dp[i][j] = matrix[i][j]
                elif matrix[i][j] == 1:
                    dp[i][j] = min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1]) + 1
                ans += dp[i][j]
        return ans

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator
Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

Which of the following uses divide and conquer strategy?


Recommended Readings

Want a Structured Path to Master System Design Too? Donā€™t Miss This!


Load More