port zebra example to cp-sat

This commit is contained in:
Laurent Perron
2018-11-19 14:13:50 -08:00
parent a5ae202fc0
commit 79a8475958

View File

@@ -33,89 +33,86 @@ Who owns a zebra and who drinks water?
"""
from __future__ import print_function
from ortools.constraint_solver import pywrapcp
from ortools.sat.python import cp_model
# pylint: disable=too-many-statements
def main():
# Create the solver.
solver = pywrapcp.Solver('zebra')
def solve_zebra():
# Create the model.
model = cp_model.CpModel()
red = solver.IntVar(1, 5, 'red')
green = solver.IntVar(1, 5, 'green')
yellow = solver.IntVar(1, 5, 'yellow')
blue = solver.IntVar(1, 5, 'blue')
ivory = solver.IntVar(1, 5, 'ivory')
red = model.NewIntVar(1, 5, 'red')
green = model.NewIntVar(1, 5, 'green')
yellow = model.NewIntVar(1, 5, 'yellow')
blue = model.NewIntVar(1, 5, 'blue')
ivory = model.NewIntVar(1, 5, 'ivory')
englishman = solver.IntVar(1, 5, 'englishman')
spaniard = solver.IntVar(1, 5, 'spaniard')
japanese = solver.IntVar(1, 5, 'japanese')
ukrainian = solver.IntVar(1, 5, 'ukrainian')
norwegian = solver.IntVar(1, 5, 'norwegian')
englishman = model.NewIntVar(1, 5, 'englishman')
spaniard = model.NewIntVar(1, 5, 'spaniard')
japanese = model.NewIntVar(1, 5, 'japanese')
ukrainian = model.NewIntVar(1, 5, 'ukrainian')
norwegian = model.NewIntVar(1, 5, 'norwegian')
dog = solver.IntVar(1, 5, 'dog')
snails = solver.IntVar(1, 5, 'snails')
fox = solver.IntVar(1, 5, 'fox')
zebra = solver.IntVar(1, 5, 'zebra')
horse = solver.IntVar(1, 5, 'horse')
dog = model.NewIntVar(1, 5, 'dog')
snails = model.NewIntVar(1, 5, 'snails')
fox = model.NewIntVar(1, 5, 'fox')
zebra = model.NewIntVar(1, 5, 'zebra')
horse = model.NewIntVar(1, 5, 'horse')
tea = solver.IntVar(1, 5, 'tea')
coffee = solver.IntVar(1, 5, 'coffee')
water = solver.IntVar(1, 5, 'water')
milk = solver.IntVar(1, 5, 'milk')
fruit_juice = solver.IntVar(1, 5, 'fruit juice')
tea = model.NewIntVar(1, 5, 'tea')
coffee = model.NewIntVar(1, 5, 'coffee')
water = model.NewIntVar(1, 5, 'water')
milk = model.NewIntVar(1, 5, 'milk')
fruit_juice = model.NewIntVar(1, 5, 'fruit juice')
old_gold = solver.IntVar(1, 5, 'old gold')
kools = solver.IntVar(1, 5, 'kools')
chesterfields = solver.IntVar(1, 5, 'chesterfields')
lucky_strike = solver.IntVar(1, 5, 'lucky strike')
parliaments = solver.IntVar(1, 5, 'parliaments')
old_gold = model.NewIntVar(1, 5, 'old gold')
kools = model.NewIntVar(1, 5, 'kools')
chesterfields = model.NewIntVar(1, 5, 'chesterfields')
lucky_strike = model.NewIntVar(1, 5, 'lucky strike')
parliaments = model.NewIntVar(1, 5, 'parliaments')
solver.Add(solver.AllDifferent([red, green, yellow, blue, ivory]))
solver.Add(
solver.AllDifferent(
[englishman, spaniard, japanese, ukrainian, norwegian]))
solver.Add(solver.AllDifferent([dog, snails, fox, zebra, horse]))
solver.Add(solver.AllDifferent([tea, coffee, water, milk, fruit_juice]))
solver.Add(
solver.AllDifferent(
[parliaments, kools, chesterfields, lucky_strike, old_gold]))
model.AddAllDifferent([red, green, yellow, blue, ivory])
model.AddAllDifferent([englishman, spaniard, japanese, ukrainian, norwegian])
model.AddAllDifferent([dog, snails, fox, zebra, horse])
model.AddAllDifferent([tea, coffee, water, milk, fruit_juice])
model.AddAllDifferent(
[parliaments, kools, chesterfields, lucky_strike, old_gold])
solver.Add(englishman == red)
solver.Add(spaniard == dog)
solver.Add(coffee == green)
solver.Add(ukrainian == tea)
solver.Add(green == ivory + 1)
solver.Add(old_gold == snails)
solver.Add(kools == yellow)
solver.Add(milk == 3)
solver.Add(norwegian == 1)
solver.Add(abs(fox - chesterfields) == 1)
solver.Add(abs(horse - kools) == 1)
solver.Add(lucky_strike == fruit_juice)
solver.Add(japanese == parliaments)
solver.Add(abs(norwegian - blue) == 1)
model.Add(englishman == red)
model.Add(spaniard == dog)
model.Add(coffee == green)
model.Add(ukrainian == tea)
model.Add(green == ivory + 1)
model.Add(old_gold == snails)
model.Add(kools == yellow)
model.Add(milk == 3)
model.Add(norwegian == 1)
all_vars = [
parliaments, kools, chesterfields, lucky_strike, old_gold, englishman,
spaniard, japanese, ukrainian, norwegian, dog, snails, fox, zebra,
horse, tea, coffee, water, milk, fruit_juice, red, green, yellow, blue,
ivory
]
distance_fox_chesterfields = model.NewIntVar(-4, 4, '')
model.Add(distance_fox_chesterfields == fox - chesterfields)
model.AddAbsEquality(1, distance_fox_chesterfields)
distance_horse_kools = model.NewIntVar(-4, 4, '')
model.Add(distance_horse_kools == horse - kools)
model.AddAbsEquality(1, distance_horse_kools)
model.Add(lucky_strike == fruit_juice)
model.Add(japanese == parliaments)
distance_norwegian_blue = model.NewIntVar(-4, 4, '')
model.Add(distance_norwegian_blue == norwegian - blue)
model.AddAbsEquality(1, distance_norwegian_blue)
solver.NewSearch(
solver.Phase(all_vars, solver.INT_VAR_DEFAULT,
solver.INT_VALUE_DEFAULT))
if solver.NextSolution():
# Solve and print out the solution.
solver = cp_model.CpSolver()
status = solver.Solve(model)
if status == cp_model.FEASIBLE:
people = [englishman, spaniard, japanese, ukrainian, norwegian]
water_drinker = [p for p in people if p.Value() == water.Value()][0]
zebra_owner = [p for p in people if p.Value() == zebra.Value()][0]
water_drinker = [p for p in people if solver.Value(p) == solver.Value(water)][0]
zebra_owner = [p for p in people if solver.Value(p) == solver.Value(zebra)][0]
print('The', water_drinker.Name(), 'drinks water.')
print('The', zebra_owner.Name(), 'owns the zebra.')
else:
print('No solutions to the zebra problem, this is unusual!')
solver.EndSearch()
if __name__ == '__main__':
main()
solve_zebra()