Magic square complete solver

magicsquare_solver.py

'''

Complete solver for all magic square 
with preset

Created on Sep 28, 2012
@author: ludoo

'''
from constraint import *
import sys

def main(show=True, preset=True):
    
    '''
    Sample config for square 3x3 and 4x4
    '''
    size=4
    if size==4:
        #http://quasistoic.org/fun/magicsquare/all.html
        rowsum = 34
        #924 possibilities
        initValue = [[1, 2, 15, 16], 
                     [13, 0, 0, 0], 
                     [0, 0, 0, 0], 
                     [0, 0, 0, 0]] 

    #http://www.mathematische-basteleien.de/magsquare.htm#The%20Magic%203x3%20Square
    if size == 3:
        rowsum = 15
        initValue = [[8,1, 6],
                     [0, 0, 0],
                     [0, 0, 0]]
        #8 possibilities
    
    problem = Problem()
    problem.addVariables(range(0, size*size), range(1, size*size+1))
    problem.addConstraint(AllDifferentConstraint())
    
    problem.addConstraint(ExactSumConstraint(rowsum), [i*(size+1) for i in range(0,size)])
    problem.addConstraint(ExactSumConstraint(rowsum), [(i+1)*(size-1) for i in range(0,size)])
    for r in range(size):
        problem.addConstraint(ExactSumConstraint(rowsum),
                              [r*size+i for i in range(size)])
        problem.addConstraint(ExactSumConstraint(rowsum),
                              [r+size*i for i in range(size)])  

    # preset.  
    if preset:
        for i in range(size) :
            for j in range(size):
                if initValue[i][j] !=0 :
                    print "var %d = %d" %( (i*10+j) , initValue[i][j])
                    problem.addConstraint(lambda var, val=initValue[i][j]:
                                          var==val, (i*size+j,))
    solutions = problem.getSolutions()
    print "Found %d solution(s)!" % len(solutions)
    if show:
        for solution in solutions:
            showSolution(solution, size)
          
def showSolution(solution, size):
    for row in range(size):
        for col in range(size):
            sys.stdout.write("%d " % solution[row*size+col])
        sys.stdout.write("\n")   
    sys.stdout.write("\n")   

if __name__ == "__main__":
    show = False
    if len(sys.argv) == 2 and sys.argv[1] == "-s":
        show = True
    elif len(sys.argv) != 1:
        sys.exit("Usage: magic.py [-s]")
    main(show)