<p>Given a <code>m * n</code> matrix of ones and zeros, return how many <strong>square</strong> submatrices have all ones.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> matrix =
[
&nbsp; [0,1,1,1],
&nbsp; [1,1,1,1],
&nbsp; [0,1,1,1]
]
<strong>Output:</strong> 15
<strong>Explanation:</strong> 
There are <strong>10</strong> squares of side 1.
There are <strong>4</strong> squares of side 2.
There is  <strong>1</strong> square of side 3.
Total number of squares = 10 + 4 + 1 = <strong>15</strong>.
</pre>

<p><strong class="example">Example 2:</strong></p>

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

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= arr.length&nbsp;&lt;= 300</code></li>
	<li><code>1 &lt;= arr[0].length&nbsp;&lt;= 300</code></li>
	<li><code>0 &lt;= arr[i][j] &lt;= 1</code></li>
</ul>

def square_is_all_ones(i, j, h, w, side_len, grid, all_ones_cols, all_ones_rows, memo):
    key = (i,j,side_len) 
    if key in memo:
        return memo[key]
    if i + side_len > h:
        memo[key] = False
        return False
    if j + side_len > w:
        memo[key] = False
        return False
    for ii in range(i, i+side_len):
        if all_ones_rows[ii]:
            continue
        for jj in range(j, j+ side_len):
            if grid[ii][jj] == 0:
                memo[key] = False
                return False
    memo[key] = True
    return True

def memoize_subsquares(memo, i, j, side_len):
    if (i,j,side_len) in memo:
        return
    for ii in range(i, i + side_len):
        iii = i + side_len - ii
        for jj in range(j, j + side_len):
            jjj = j + side_len - jj
            for ss in range(min(iii, jjj)):
                key = (ii, jj, ss)
                memo[key] = True

class Solution:
    def countSquares(self, matrix: List[List[int]]) -> int:
        memo = {}
        h = len(matrix)
        w = len(matrix[0])
        
        all_ones_cols = [True] * w
        all_ones_rows = [True] * h
        for i in range(h):
            for j in range(w):
                if matrix[i][j] == 0:
                    all_ones_cols[j] = False
                    all_ones_rows[i] = False
        # print(all_ones_cols)
        # print(all_ones_rows)
        count = 0
        
        sides_to_see = min(h,w)+1
        # print(sides_to_see)
        for i in range(h):
            for j in range(w):
                for side in range(1, sides_to_see):
                    if square_is_all_ones(i, j, h, w, side, matrix, all_ones_cols, all_ones_rows, memo):
                        # print(f"i:{i}, j:{j}, side:{side}")
                        count += 1
                        memoize_subsquares(memo, i, j, side)
                    else:
                        break
                        
        return count
35600