'We solve the Zebra Puzzle with the help of a MIP model and SCIP

zebra.py

'''We solve the Zebra Puzzle with the help of a MIP model and SCIP'''

from zibopt import scip

# see https://en.wikipedia.org/wiki/Zebra_Puzzle
houses = range(5)
nations = ['England', 'Japan', 'Norway', 'Spain',  'Ukraine']
pets = ['dog', 'fox', 'horse', 'snails', 'zebra']
drinks = ['coffee', 'milk', 'orange juice', 'tea', 'water']
colors = ['blue', 'green', 'ivory', 'red', 'yellow']
cigarettes = ['Chesterfield', 'Kool', 'Lucky Strike', 'Old Gold', 'Parliament']

s = scip.solver()

# create assignment variables for every pair of item, house
a = dict()
for item in sum([nations, pets, drinks, colors, cigarettes], []):
    a[item] = [ s.variable(scip.BINARY) for _ in houses ]

def house_sum(item):
    '''Express the actual house number in terms of the assignment'''
    return sum(house * a[item][house] for house in houses)

# add assignment constraints:
# every house has exactly one item of every category, and vice versa
for cat in [nations, pets, drinks, colors, cigarettes]:
    for house in houses:
        s += sum(a[item][house] for item in cat) == 1.0
    for item in cat:
        s += sum(a[item][house] for house in houses) == 1.0

# add rest of the constraints:
#
# (There are five houses.)
# The Englishman lives in the red house.
s += house_sum('England') == house_sum('red')

# The Spaniard owns the dog.
s += house_sum('Spain') == house_sum('dog')

# Coffee is drunk in the green house.
s += house_sum('coffee') == house_sum('green')

# The Ukrainian drinks tea.
s += house_sum('Ukraine') == house_sum('tea')

# The green house is immediately to the right of the ivory house.
s += house_sum('green') == house_sum('ivory') + 1

# The Old Gold smoker owns snails.
s += house_sum('Old Gold') == house_sum('snails')

# Kools are smoked in the yellow house.
s += house_sum('Kool') == house_sum('yellow')

# Milk is drunk in the middle house.
s += a['milk'][2] == 1

# The Norwegian lives in the first house.
s += a['Norway'][0] == 1

# The man who smokes Chesterfields lives in the house next to the
# man with the fox.
helper = s.variable(scip.BINARY)
s += house_sum('Chesterfield') - house_sum('fox') == 2 * helper - 1

# Kools are smoked in the house next to the house where the horse is
# kept. [should be "... a house ...", see discussion below]
helper = s.variable(scip.BINARY)
s += house_sum('Kool') - house_sum('horse') == 2 * helper - 1

# The Lucky Strike smoker drinks orange juice.
s += house_sum('Lucky Strike') == house_sum('orange juice')

# The Japanese smokes Parliaments.
s += house_sum('Japan') == house_sum('Parliament')

# The Norwegian lives next to the blue house.
helper = s.variable(scip.BINARY)
s += house_sum('Norway') - house_sum('blue') == 2 * helper - 1

# In the interest of clarity, it must be added that each of the five
# houses is painted a different color, and their inhabitants are of
# different national extractions, own different pets, drink different
# beverages and smoke different brands of American cigarets [sic]. One
# other thing: in statement 6, right means your right.

sol = s.maximize()

def house_sol(item):
    '''Get the house number from the solution'''
    house = sum(house * sol[ a[item][house] ] for house in houses)
    return '%s is at house %d' % (item, round(house))

def item_sol(cat, house):
    '''Get the item of the category at some house'''
    res = [item for item in cat if sol[ a[item][house] ] > 0.5]
    assert len(res) == 1
    return res[0]

# Now, who drinks water?
print(house_sol('water'))

# Who owns the zebra?
print(house_sol('zebra'))
print()

# all results in tabular form
for cat in [colors, nations, drinks, cigarettes, pets]:
    print(' | '.join('%-12s' % item_sol(cat, house) for house in houses))