EinsteinPuzzle.py

import itertools
nationalities = ["Brit", "Dane", "Norwegian", "German", "Swede"]
# assign arbitrary number to each nationality
Brit, Dane, Norwegian, German, Swede = range(len(nationalities))
# this returns an iterater over all the permutations of 5 numbers
choices = lambda: itertools.permutations(range(len(nationalities)))

next_to_ltor = lambda l, r: l - 1 == r
next_to = lambda l, r: abs(l - r) == 1


def get_fish():
    """
    Slow version.
    Checks through all 5!^5 permutations, using the constraints from the
    puzzle in order.
    """
    answer = [fish
              for red, yellow, green, blue, white in choices()
              for milk, water, beer, coffee, tea in choices()
              for houses in choices()
              for blends, blue_master, prince, dunhill, pall_mall in choices()
              for dogs, cats, birds, horses, fish in choices()
              # The Brit lives in the Red house.
              if red == Brit
              # The Swede keeps dogs as pets.
              if Swede == dogs
              # The Dane drinks tea.
              if Dane == tea
              # The Green house is next to the White house, on the left.
              if next_to_ltor(houses[green], houses[white])
              # The owner of the Green house drinks coffee.
              if green == coffee
              # The person who smokes Pall Mall rears birds.
              if pall_mall == birds
              # The owner of the Yellow house smokes Dunhill.
              if yellow == dunhill
              # The man living in the centre house drinks milk.
              if houses[2] == milk
              # The Norwegian lives in the first house.
              if houses[0] == Norwegian
              # The man who smokes Blends lives next to the one who keeps
              # cats.
              if next_to(houses[blends], houses[cats])
              # The man who keeps horses lives next to the man who smokes
              # Dunhill.
              if next_to(houses[horses], houses[dunhill])
              # The man who smokes Blue Master drinks beer.
              if blue_master == beer
              # The German smokes Prince.
              if German == prince
              # The Norwegian lives next to the Blue house.
              if next_to(houses[Norwegian], houses[blue])
              # The man who smokes Blends has a neighbour who drinks water.
              if next_to(houses[blends], houses[water])
              ]
    assert len(answer) == 1
    return nationalities[answer[0]]


def get_fish_optimized():
    """
    Optimized version.
    Fails early for invalid combinations before even considering later ones.
    """
    answer = [fish
              for red, yellow, green, blue, white in choices()
              # The Brit lives in the Red house.
              if red == Brit
              for milk, water, beer, coffee, tea in choices()
              # The Dane drinks tea.
              if Dane == tea
              # The owner of the Green house drinks coffee.
              if green == coffee
              for houses in choices()
              # The Norwegian lives in the first house.
              if houses[0] == Norwegian
              # The Green house is next to the White house, on the left.
              if next_to_ltor(houses[green], houses[white])
              # The man living in the centre house drinks milk.
              if houses[2] == milk
              # The Norwegian lives next to the Blue house.
              if next_to(houses[Norwegian], houses[blue])
              for blends, blue_master, prince, dunhill, pall_mall in choices()
              # The German smokes Prince.
              if German == prince
              # The owner of the Yellow house smokes Dunhill.
              if yellow == dunhill
              # The man who smokes Blue Master drinks beer.
              if blue_master == beer
              # The man who smokes Blends has a neighbour who drinks water.
              if next_to(houses[blends], houses[water])
              for dogs, cats, birds, horses, fish in choices()
              # The Swede keeps dogs as pets.
              if Swede == dogs
              # The person who smokes Pall Mall rears birds.
              if pall_mall == birds
              # The man who smokes Blends lives next to the one who keeps
              # cats.
              if next_to(houses[blends], houses[cats])
              # The man who keeps horses lives next to the man who smokes
              # Dunhill.
              if next_to(houses[horses], houses[dunhill])
              ]
    assert len(answer) == 1
    return nationalities[answer[0]]

print(get_fish_optimized())