Grid problem solver with Z3

grid_coloring_z3.py

from z3 import *
import itertools

# A color type
num_colors = 8
Color, color_consts = EnumSort( 'Color',
                                [ f"{c}" for c in range(num_colors) ] )

# Edges
#  --0,0-- --0,1-- --0,2--  horizontal
# |       |
# 0,0    0,1 vertical
# |       |
#  --1,0--
#
# Grid, not torus, so left and right sides do not coincide

x_tiles = 10
y_tiles = 7
# 8 choose 4 = 70
# 70 * 4 = 280
# 8 * 7 = 56 ordered pairs of colors
target = 5

# 6 choose 4 = 15
# 15 * 4 = 60 corners
# 6 * 5 = 30 ordered pairs of colors
#x_tiles=5
#y_tiles=3
#target = 2

horizontal = [ [ Const(f"h_{i}_{j}", Color) for j in range( x_tiles ) ]
                for i in range( y_tiles + 1) ]
vertical = [ [ Const(f"v_{i}_{j}", Color) for j in range( x_tiles + 1 ) ]
                 for i in range( y_tiles ) ]

s = Solver()

def tile( y, x ):
    return [horizontal[y][x],
        horizontal[y+1][x],
        vertical[y][x],
        vertical[y][x+1]]

def make_tile( y, x ):
    t = EmptySet(Color)
    for e in tile( y, x ):
        t = SetAdd( t, e )
    return t

tiles = [make_tile(y,x) for x in range( x_tiles ) for y in range( y_tiles )]

# Maybe this is cheaper than a distinctness check-- we know what the tiles
# must be if they are all distinct.
for four_colors in itertools.combinations( color_consts, r=4):
    desired = EmptySet(Color)
    for c in four_colors:
        desired = SetAdd( desired, c )
    s.add( Or( *[ desired == t for t in tiles] ) )

def corners( y, x ):
    edges = tile( y, x )
    return zip( edges, edges[1:] + [edges[0]] )

all_corners = []
for y in range( y_tiles ):
    for x in range( x_tiles ):
        for c in corners(y,x):
            all_corners.append( c )

for two_colors in itertools.combinations( color_consts, r=2 ):
    clause = []
    for c in all_corners:
        clause.append( And( c[0] == two_colors[0], c[1] == two_colors[1] ) )
    s.add( AtLeast( *( clause + [target] ) ) )
    # Reverse order
    clause = []
    for c in all_corners:
        clause.append( And( c[0] == two_colors[1], c[1] == two_colors[0] ) )
    s.add( AtLeast( *( clause + [target] ) ) )
        
def show( model ):
    for i in range( y_tiles + 1 ):
        for j in range( x_tiles ):
            print( " --{}--".format( model[horizontal[i][j]] ),
                  end="" )
        print()
        if i != y_tiles:
            for j in range( x_tiles+1 ):
                print( "|     ", end="" )
            print()
            for j in range( x_tiles+1 ):
                print( "{}     ".format( model[vertical[i][j]] ), end="" )
            print()
            for j in range( x_tiles+1 ):
                print( "|     ", end="" )
            print()
    count_by_pair = {}
    for c in all_corners:
        key = (str(model[c[0]]),
               str(model[c[1]]))
        count_by_pair[key] = count_by_pair.get(key, 0)+ 1
    print( count_by_pair )

grid_coloring_z3_alt.py

from z3 import *
import itertools

# A color type
num_colors = 8
Color, color_consts = EnumSort( 'Color',
                                [ f"{c}" for c in range(num_colors) ] )

# Each tile is one of 8C4 colors; the question is which order the
# colors appear around the edge,  so there are 24 possibilities for each.
# The tiles can be in one of 70 locations.
# We'll model this as a mapping from tile to orientation.

if num_colors == 8:
    x_tiles = 10
    y_tiles = 7
    # 8 choose 4 = 70
    # 70 * 4 = 280
    # 8 * 7 = 56 ordered pairs of colors
    # target = 5
    target = 0

if num_colors == 6:
    # 6 choose 4 = 15
    #15 * 4 = 60 corners
    # 6 * 5 = 30 ordered pairs of colors
    x_tiles=5
    y_tiles=3
    target = 2
    #target = 1

if num_colors == 5:
    # 5 choose 4 = 5
    # 5 * 4 = 20 corners
    # 5 * 4 = 20 ordered pairs of colors
    x_tiles = 5
    y_tiles = 1
    target = 1

class Tile(object):
    def __init__( self, four_colors, solver ):
        label = "".join( str(c) for c in four_colors )
        self.left = Const( label + "-l", Color )
        self.right = Const( label + "-r", Color )
        self.top = Const( label + "-t", Color )
        self.bottom = Const( label + "-b", Color )

        # Possible color assignments
        self.orderings = [ And( self.left == a, self.right == b,
                                self.top == c, self.bottom == d )
                           for (a,b,c,d) in itertools.permutations( four_colors ) ]
        solver.add( Or( self.orderings ) )

        # Possible positions
        self.x = Int( f"{label}-x" )
        self.y = Int( f"{label}-y" )
        solver.add( And( self.x >= 0, self.x < x_tiles ) )
        solver.add( And( self.y >= 0, self.y < y_tiles ) )

        self.corners = [ (self.left, self.top),
                         (self.top, self.right),
                         (self.right, self.bottom),
                         (self.bottom, self.left ) ]
                         

s = Solver()
tiles = [ Tile( four_colors, s ) for four_colors in itertools.combinations( color_consts, r=4 ) ]

for t1, t2 in itertools.combinations( tiles, r=2 ):
    # No two tiles at the same position
    s.add( Not( And( t1.x == t2.x, t1.y == t2.y ) ) )
    
    # If two tiles are adjacent, then their corresponding colors must match.
    s.add( Implies( And( t1.x + 1 == t2.x, t1.y == t2.y ),
                    t1.right == t2.left ) )
    s.add( Implies( And( t1.x == t2.x + 1, t1.y == t2.y ),
                    t1.left == t2.right) )
    s.add( Implies( And( t1.x == t2.x, t1.y + 1 == t2.y ),
                    t1.bottom == t2.top ) )
    s.add( Implies( And( t1.x == t2.x, t1.y == t2.y + 1),
                    t1.top == t2.bottom ) )


# At least "target" representative of each color pair, represented by Corner objects
class Corner(object):
     def __init__( self, two_colors, k, solver ):
        self.colors = two_colors
        label = "corner-{}{}-{}".format( two_colors[0], two_colors[1], k )
        self.x = Int( f"{label}-x" )
        self.y = Int( f"{label}-y" )
        self.c = Int( f"{label}-c" )
        solver.add( And( self.x >= 0, self.x < x_tiles ) )
        solver.add( And( self.y >= 0, self.y < y_tiles ) )
        solver.add( And( self.c >= 0, self.c < 4 ) )

corners = []
corners_by_color = {}

for two_colors in itertools.combinations( color_consts, r=2 ):
    for k in range( target ):
        for tc in [ tuple(two_colors), (two_colors[1], two_colors[0]) ]:
            c = Corner( tc, k, s )
            corners.append( c )
            corners_by_color[tc] = corners_by_color.get( tc, [] ) + [c] 

# No two corners in the same position
for c1, c2 in itertools.combinations( corners, r=2 ):
    s.add( Not( And( c1.x == c2.x, c1.y == c2.y, c1.c == c2.c ) ) )


# Corners must match the tile colors
for t in tiles:
    for c in corners:
        s.add( Implies( And( t.x == c.x, t.y == c.y, c.c == 0 ),
                        And( t.corners[0][0] == c.colors[0], t.corners[0][1] == c.colors[1] ) ) )
        s.add( Implies( And( t.x == c.x, t.y == c.y, c.c == 1 ),
                        And( t.corners[1][0] == c.colors[0], t.corners[1][1] == c.colors[1] ) ) )
        s.add( Implies( And( t.x == c.x, t.y == c.y, c.c == 2 ),
                        And( t.corners[2][0] == c.colors[0], t.corners[2][1] == c.colors[1] ) ) )
        s.add( Implies( And( t.x == c.x, t.y == c.y, c.c == 3 ),
                        And( t.corners[3][0] == c.colors[0], t.corners[3][1] == c.colors[1] ) ) )

# Alternate conditions not based on explicit corner variables
if False:
    for two_colors in itertools.combinations( color_consts, r=2 ):
        clause = []
        
        for t in tiles:
            for c in t.corners:
                clause.append( And( c[0] == two_colors[0], c[1] == two_colors[1] ) )
        s.add( AtLeast( *( clause + [target] ) ) )
        
        # Reverse order
        clause = []
        for t in tiles:
            for c in t.corners:
                clause.append( And( c[0] == two_colors[1], c[1] == two_colors[0] ) )
        s.add( AtLeast( *( clause + [target] ) ) )
    
def show( model, showCorners = False ):
    array = [ [ None for j in range( x_tiles )] for i in range( y_tiles ) ]

    for t in tiles: 
        array[model[t.y].as_long()][model[t.x].as_long()] = t

    for i in range( y_tiles ):
        for j in range( x_tiles ):
            print( " -----", end="" )
        print()

        for j in range( x_tiles ):
            print( "|  {}  ".format( model[array[i][j].top] ), end="" )
        print()
        
        for j in range( x_tiles ):
            print( "|{}   {}".format( model[array[i][j].left], model[array[i][j].right] ), end="" )
        print()

        for j in range( x_tiles ):
            print( "|  {}  ".format(model[array[i][j].bottom] ), end="" )
        print()

    if showCorners:
        print()
        for tc, corners in corners_by_color.items():
            print( tc, ": ", end = "" )
            for c in corners:
                print( "{},{},{} ".format( model[c.y], model[c.x], model[c.c] ), end="" )
            print()