Count of square submatrices with average at least K | Function to count submatrixes with average greater than or equals to K ; Stores count of submatrices ; Stores the prefix sum of matrix ; Iterate over the range [ 1 , N ] ; Iterate over the range [ 1 , M ] ; Update the prefix sum ; Iterate over the range [ 1 , N ] ; Iterate over the range [ 1 , M ] ; Iterate until l and r are greater than 0 ; Update count ; Stores sum of submatrix with bottom right corner as ( i , j ) and top left corner as ( l , r ) ; If sum1 is less than or equal to sum2 ; Increment cnt by 1 ; Return cnt as the answer ; Driver Code ; Given Input ; Function Call
def cntMatrices ( arr , N , M , K ) : NEW_LINE INDENT cnt = 0 NEW_LINE pre = [ [ 0 for i in range ( M + 1 ) ] for i in range ( N + 1 ) ] NEW_LINE for i in range ( 1 , N + 1 ) : NEW_LINE INDENT for j in range ( 1 , M + 1 ) : NEW_LINE INDENT pre [ i ] [ j ] = ( arr [ i - 1 ] [ j - 1 ] + pre [ i - 1 ] [ j ] + pre [ i ] [ j - 1 ] - pre [ i - 1 ] [ j - 1 ] ) NEW_LINE DEDENT DEDENT for i in range ( 1 , N + 1 ) : NEW_LINE INDENT for j in range ( 1 , M + 1 ) : NEW_LINE INDENT l , r = i , j NEW_LINE while l > 0 and r > 0 : NEW_LINE INDENT sum1 = ( K * ( i - l + 1 ) * ( i - r + 1 ) ) NEW_LINE sum2 = ( pre [ i ] [ j ] - pre [ l - 1 ] [ r ] - pre [ l ] [ r - 1 ] + pre [ l - 1 ] [ r - 1 ] ) NEW_LINE if ( sum1 <= sum2 ) : NEW_LINE INDENT cnt += 1 NEW_LINE DEDENT l -= 1 NEW_LINE r -= 1 NEW_LINE DEDENT DEDENT DEDENT return cnt NEW_LINE DEDENT if __name__ == ' _ _ main _ _ ' : NEW_LINE INDENT arr = [ [ 2 , 2 , 3 ] , [ 3 , 4 , 5 ] , [ 4 , 5 , 5 ] ] NEW_LINE K = 4 NEW_LINE N = len ( arr ) NEW_LINE M = len ( arr [ 0 ] ) NEW_LINE print ( cntMatrices ( arr , N , M , K ) ) NEW_LINE DEDENT