Zebra Puzzle

This is a classic logic puzzle that is useful for demonstrating a constraint satisfaction problem (CSP). It is sometimes referred to as either Einstein's Puzzle, or Einstein's Riddle, due to an apocryphal story that it was invented by Albert Einstein as a boy.

From http://en.wikipedia.org/wiki/Zebra_Puzzle:

  1. There are five houses.
  2. The Englishman lives in the red house.
  3. The Spaniard owns the dog.
  4. Coffee is drunk in the green house.
  5. The Ukrainian drinks tea.
  6. The green house is immediately to the right of the ivory house.
  7. The Old Gold smoker owns snails.
  8. Kools are smoked in the yellow house.
  9. Milk is drunk in the middle house.
  10. The Norwegian lives in the first house.
  11. The man who smokes Chesterfields lives in the house next to the man with the fox.
  12. Kools are smoked in the house next to the house where the horse is kept.
  13. The Lucky Strike smoker drinks orange juice.
  14. The Japanese smokes Parliaments.
  15. The Norwegian lives next to the blue house.

Answer the following:

  1. Who drinks water?
  2. Who owns the zebra?

In [ ]:
from pyeda.inter import *

In [ ]:
nationalities = {"England", "Japan", "Norway", "Spain", "Ukraine"}
colors = {"blue", "red", "green", "ivory", "yellow"}
pets = {"dog", "fox", "horse", "snails", "zebra"}
drinks = {"coffee", "milk", "oj", "tea", "water"}
smokes = {"Chesterfield", "OldGold", "Kools", "LuckyStrike", "Parliament"}

In [ ]:
# Create Boolean variables
xs = []
for i in range(5):
    xs.append({})
    for group in (nationalities, colors, pets, drinks, smokes):
        for name in group:
            xs[i][name] = exprvar(name, i)

In [ ]:
# Start building a list of constraints
cons = []

In [ ]:
# Basic uniqueness constraints
for group in (nationalities, colors, pets, drinks, smokes):
    # Each house has exactly one {name}
    for i in range(5):
        cons.append(OneHot(*[xs[i][name] for name in group]))
    # Each {name} is only in one house
    for name in group:
        cons.append(OneHot(*[xs[i][name] for i in range(5)]))

In [ ]:
# 2. The Englishman lives in the red house
for i in range(5):
    cons.append(xs[i]["England"] >> xs[i]["red"])

In [ ]:
# 3. The Spaniard owns the dog
for i in range(5):
    cons.append(xs[i]["Spain"] >> xs[i]["dog"])

In [ ]:
# 4. Coffee is drunk in the green house
for i in range(5):
    cons.append(xs[i]["coffee"] >> xs[i]["green"])

In [ ]:
# 5. The Ukrainian drinks tea
for i in range(5):
    cons.append(xs[i]["Ukraine"] >> xs[i]["tea"])

In [ ]:
# 6. The green house is immediately to the right of the ivory house
cons += [
    ~xs[0]["green"],
    xs[0]["ivory"] >> xs[1]["green"],
    xs[1]["ivory"] >> xs[2]["green"],
    xs[2]["ivory"] >> xs[3]["green"],
    xs[3]["ivory"] >> xs[4]["green"],
    ~xs[4]["ivory"],
]

In [ ]:
# 7. The Old Gold smoker owns snails
for i in range(5):
    cons.append(xs[i]["OldGold"] >> xs[i]["snails"])

In [ ]:
# 8. Kools are smoked in the yellow house
for i in range(5):
    cons.append(xs[i]["Kools"] >> xs[i]["yellow"])

In [ ]:
# 9. Milk is drunk in the middle house
cons.append(xs[2]["milk"])

In [ ]:
# 10. The Norwegian lives in the first house
cons.append(xs[0]["Norway"])

In [ ]:
# 11. The man who smokes Chesterfields lives in the house next to the man with the fox
cons += [
    xs[0]["Chesterfield"] >> xs[1]["fox"],
    xs[1]["Chesterfield"] >> (xs[0]["fox"] | xs[2]["fox"]),
    xs[2]["Chesterfield"] >> (xs[1]["fox"] | xs[3]["fox"]),
    xs[3]["Chesterfield"] >> (xs[2]["fox"] | xs[4]["fox"]),
    xs[4]["Chesterfield"] >> xs[3]["fox"],
]

In [ ]:
# 12. Kools are smoked in the house next to the house where the horse is kept
cons += [
    xs[0]["Kools"] >> xs[1]["horse"],
    xs[1]["Kools"] >> (xs[0]["horse"] | xs[2]["horse"]),
    xs[2]["Kools"] >> (xs[1]["horse"] | xs[3]["horse"]),
    xs[3]["Kools"] >> (xs[2]["horse"] | xs[4]["horse"]),
    xs[4]["Kools"] >> xs[3]["horse"],
]

In [ ]:
# 13. The Lucky Strike smoker drinks orange juice
for i in range(5):
    cons.append(xs[i]["LuckyStrike"] >> xs[i]["oj"])

In [ ]:
# 14. The Japanese smokes Parliaments
for i in range(5):
    cons.append(xs[i]["Japan"] >> xs[i]["Parliament"])

In [ ]:
# 15. The Norwegian lives next to the blue house
cons += [
    xs[0]["Norway"] >> xs[1]["blue"],
    xs[1]["Norway"] >> (xs[0]["blue"] | xs[2]["blue"]),
    xs[2]["Norway"] >> (xs[1]["blue"] | xs[3]["blue"]),
    xs[3]["Norway"] >> (xs[2]["blue"] | xs[4]["blue"]),
    xs[4]["Norway"] >> xs[3]["blue"],
]

In [ ]:
# Merge all the constraints together
F = And(*cons).to_cnf()

In [ ]:
# The number of clauses in the CNF
len(F.xs)

In [ ]:
# Solve the puzzle
soln = F.satisfy_one()

In [ ]:
# Verify there is only one solution
assert F.satisfy_count() == 1

In [ ]:
# Rearrange the solution into a table
table = [{} for i in range(5)]
for var, val in soln.items():
    if val:
        if var.name in nationalities:
            table[var.indices[0]]["nationality"] = var.name
        elif var.name in colors:
            table[var.indices[0]]["color"] = var.name
        elif var.name in pets:
            table[var.indices[0]]["pet"] = var.name
        elif var.name in drinks:
            table[var.indices[0]]["drink"] = var.name
        elif var.name in smokes:
            table[var.indices[0]]["smoke"] = var.name

In [ ]:
# Who drinks water?
for i in range(5):
    if table[i]["drink"] == "water":
        print("The man from", table[i]["nationality"], "drinks water.")

In [ ]:
# Who owns the zebra?
for i in range(5):
    if table[i]["pet"] == "zebra":
        print("The man from", table[i]["nationality"], "owns the zebra.")

In [ ]:
# The full solution
table

In [ ]: