solve-puzzle.py

#!/usr/bin/env python
"""solve-puzzle.py: Solve Sergey's puzzle.

Usage:

    $ wget http://codespeak.net/svn/user/niemeyer/constraint/trunk/constraint.py
    $ python solve-puzzle.py [matrix size (should be perfect square)]

Example:

    $ ./solve-puzzle.py 4

    |3 1|0 2|
    |0 2|3 1|
    |2 0|1 3|
    |1 3|2 0|
    
    |3 0|1 2|
    |2 1|0 3|
    |0 3|2 1|
    |1 2|3 0|
    
       ...
"""
import math, sys
from operator import itemgetter
from constraint import AllDifferentConstraint, MinConflictsSolver, Problem

try: N = int(sys.argv[1])
except IndexError:
    N = 4 # NxN matrix
M = int(math.sqrt(N) + .5) # block size (M*M blocks in the matrix)
assert M*M == N
R = range(N)
# each cell in the matrix is a variable (name of the variable is its number)
v = [[N*i + j for j in R] for i in R]


def get_block_indexes(i, cond=lambda x,y: True):
    """Yields indexes that belong to `i'th matrix block."""
    top, left = divmod(i, M)
    # e.g., 7th block (zero-based): top, left = 2, 1 -> (j+6, k+3)
    return ((j+M*top, k+M*left) for j in range(M) for k in range(M)
            if cond(j, k))
    
    
    
problem = Problem()
# variables values are in [0, N)
problem.addVariables(sum(v, []), R)
# add constraints
add_adc = lambda vars_: problem.addConstraint(AllDifferentConstraint(), vars_)
for i in R:
    # all variables in the `i'th row are different
    add_adc([v[i][j] for j in R])
    # all variables in the `i'th column are different
    add_adc([v[j][i] for j in R])
    # all variables in the `i'th block are different
    # (there are MxM blocks in the matrix)
    add_adc([v[j][k] for j, k in get_block_indexes(i)])
    # all variable on the main blocks' diagonals are different
    add_adc([v[j][k] for j, k in get_block_indexes(i, lambda x,y: x == y)])
    add_adc([v[j][k] for j, k in get_block_indexes(i, lambda x,y: x == (M - 1 - y))])
# all variables on the main diagonals are different
add_adc([v[i][j] for i in R for j in R if i == j])
add_adc([v[i][j] for i in R for j in R if i == (N - 1 - j)])


def print_solution(solution):
    # solution: varname -> value
    w = sys.stdout.write
    for i in R:
        for j in R:
            if j % M == 0: w('|')
            else: w(' ')
            w(str(solution[v[i][j]]))
        print '|'
    print


# find & print solutions
try: solutions = problem.getSolutionIter()
except NotImplementedError:
    try: solutions = problem.getSolutions()
    except NotImplementedError:
        solutions = [problem.getSolution()]

for i, solution in enumerate(solutions):
    if solution is None:
        i = -1
        break
    print_solution(solution)
    if i == 10:
        break
print i+1