find all solutions of n-order magic square

magic_square.py

# -*- coding: utf-8 -*-

"""
find all solutions of n-order magic square
"""

import itertools
 
class Solver(object):
    def __init__(self, n, square=None, magic_sum=None):
        assert n > 0
        self.order = n
        self.order_square = n*n

        self.max_num = n*n
        self.magic_sum = n * (n*n + 1)/2

        if magic_sum:
            assert magic_sum >= self.magic_sum
            self.max_num = magic_sum - 3
            self.magic_sum = magic_sum

        self.remains = set(range(1, self.max_num+1))

        if square:
            assert len(square) == n
            assert map(len, square) == [n] * n
            self.square = square
            self.remains -= set(a for a in itertools.chain(*self.square) if a > 0)
        else:
            self.square = []
            for x in xrange(n):
                self.square.append([0] * n)

        print self.order
        print self.max_num, self.magic_sum
        print self.remains
        print self.square
        print '---'

        self.s_count = 0



    def _push(self, i, j, v):
        assert v in self.remains
        self.square[i][j] = v
        self.remains.remove(v)

    def _pop(self, i, j):
        v = self.square[i][j]
        assert v and v not in self.remains
        self.remains.add(v)
        self.square[i][j] = 0
 
 
    def _pos2ij(self, pos):
        return pos / self.order, pos % self.order
 
    def _is_ok(self, i, j):
        if j == self.order-1:
            if sum(self.square[i]) != self.magic_sum:
                return False
        if i == self.order-1:
            s = 0
            for ii in xrange(self.order):
                s += self.square[ii][j]
            if s != self.magic_sum:
                return False
     
        if i == self.order-1 and j == self.order-1:
            s = 0
            for ii in xrange(self.order):
                s += self.square[ii][ii]
            if s != self.magic_sum:
                return False
            
        if i == self.order-1 and j == 0:
            s = 0
            for ii in xrange(self.order):
                s += self.square[ii][self.order-1-ii]
            if s != self.magic_sum:
                return False
     
        return True
 
 
    def run(self, pos=0):
        if pos >= self.order_square:
            self.s_count +=1
            print '--Bingo!~', self.s_count, self.square
            return
     
        i, j = self._pos2ij(pos)
     
        if self.square[i][j] > 0:
            self.run(pos+1)
            return
        
        for v in list(self.remains):
            self._push(i, j, v)
            if self._is_ok(i, j):
                self.run(pos+1)
            self._pop(i, j)

 
if __name__ == '__main__':
    #s = Solver(1)
    #s = Solver(1, max_num=100, magic_sum=100)

    #s = Solver(2)
    #s = Solver(2, max_num=100, magic_sum=150)

    #square = [[0, 0, 0], [9, 0, 0], [0, 0, 0]]
    #s = Solver(3, square)
    #s = Solver(3, magic_sum=23) # solutions: 0
    #s = Solver(3, magic_sum=18) # solutions: 24
    #s = Solver(3, magic_sum=34) # solutions: 0
    #s = Solver(3, magic_sum=20) # solutions: 0

    #square = [ [2, 3, 0, 0], [4, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0] ]
    #s = Solver(4, square)

    s = Solver(4)

    s.run()