Magic square searcher

gistfile1.py

# -*- coding: utf-8 -*-
"""
Created on Sat Apr 18 23:52:19 2015

@author: euphorbium

Finds magic squares in two dimmensional matrix
"""
import itertools
import numpy as np


def magic_squares(matrix):
    found = 0
    max_square = min(matrix.shape)
    for size in xrange(2, max_square+1):
        for i in xrange(matrix.shape[0]-size+1):
            for j in xrange(matrix.shape[1]-size+1):
                square = matrix[i:i+size, j:j+size]
                if len(set(map(
                sum, itertools.chain(square, 
                                     square.transpose(),
                                     [square.diagonal()],
                                     [square[::-1].diagonal()]))))==1:
                    found+=1
    return found
    
if __name__=='__main__':
    matrix = np.array([[4, 9, 2, 3, 5],
           [3, 5, 7, 4, 2],
           [8, 1, 6, 6, 2],
           [1, 1, 6, 6, 2]])
    
    print magic_squares(matrix)