Problem statement

https://leetcode.com/problems/magic-squares-in-grid/

Solution

Just check every possible 3x3 squares: there will be O(mn) of them. We need to evaluate sums of rows, columns, diagonals and check that we have numbers from 1 to 9.

Complexity

Time complexity is O(9mn), space is O(9).

Code

class Solution:
    def numMagicSquaresInside(self, grid):
        m, n = len(grid), len(grid[0])
        ans = 0
        for i in range(0, m-2):
            for j in range(0, n-2):
                matrix = [grid[q][j:j+3] for q in range(i,i+3)]
                elements = set(chain(*matrix))
                cols = [sum(i) for i in zip(*matrix)]
                rows = [sum(i) for i in matrix]
                d1 = [sum(matrix[i][i] for i in range(3))]
                d2 = [sum(matrix[i][2-i] for i in range(3))]
                if len(set(cols + rows + d1 + d2)) == 1 and elements == set(range(1,10)):
                    ans += 1
                    
        return ans