HTMLify

class Solution:
    def countSquare(self, mat, x):
        n = len(mat)
        m = len(mat[0])
        
        pref = [[0] * (m + 1) for _ in range(n + 1)]
        for i in range(1, n + 1):
            for j in range(1, m + 1):
                pref[i][j] = (mat[i-1][j-1] + 
                             pref[i-1][j] + 
                             pref[i][j-1] - 
                             pref[i-1][j-1])
        
        count = 0
        for i in range(n):
            for j in range(m):
                for s in range(1, min(n - i, m - j) + 1):
                    r2, c2 = i + s, j + s

                    current_sum = (pref[r2][c2] - 
                                  pref[i][c2] - 
                                  pref[r2][j] + 
                                  pref[i][j])
                    
                    if current_sum == x:
                        count += 1
                        
        return count