Zebra puzzle / Einstein's riddle

zebrapuzzle.py

# This script is a solution to this version of the Zebra Puzzle:
# http://csl.name/post/einsteins-puzzle/
# About which everything can be found at:
# http://en.wikipedia.org/wiki/Zebra_Puzzle

# After trying to generate all possible compositions of houses, nationality, color, drink, smoke and pet it still took too long for my machine to find the solution. So I took some business rules and applied these before generating the compositions. Eventually it still takes around a minute or two to find and print the solution. It was a nice challenge for a beginning programmer.

# Business rules:
#  1 The Brit lives in a red house.
#  2 The Swede keeps dogs as pets.
#  3 The Dane drinks tea.
#  4 The green house is on the left of the white, next to it.
#  5 The green house owner drinks coffee.
#  6 The person who smokes Pall Mall rears birds.
#  7 The owner of the yellow house smokes Dunhill.
#  8 The man living in the house right in the center drinks milk.
#  9 The Norwegian lives in the first house.
# 10 The man who smokes blend lives next to the one who keeps cats.
# 11 The man who keeps horses lives next to the man who smokes Dunhill.
# 12 The owner who smokes Blue Master drinks beer.
# 13 The German smokes Prince.
# 14 The Norwegian lives next to the blue house.
# 15 The man who smokes blend has a neighour who drinks water.

# itertools is required to arrange for the different permutations in order to create the different compositions.
import itertools

number = [1,2,3,4,5]
countries = ["UK","S","DK","DE"]
colors = ["yellow","blue","red","green","white"]
drinks = ["beer","water","coffee","tea","milk"]
smokes = ["Pall Mall","blend","Prince","Dunhill","Blue Master"]
pets = ["fish","dogs","horse","cats","birds"]

#  9 The Norwegian lives in the first house.
# Inputting (countries,"NO") makes sure only composition where this business rule applies are compositions.
def countryFirst(input,country):
  permutations = list(itertools.permutations(input))
  output = []
  for n in range(0,len(permutations)):
    converted = list(permutations[n])
    converted.insert(0,country)
    output.append(converted)
  print "Remaining items in countries " + str(len(output))
  return output

#  4 The green house is on the left of the white, next to it.
# 14 The Norwegian lives next to the blue house.
# Inputting (colors) makes sure white house is never first, the green house is never last and the blue house is in position no. 2
def colorConditions(input):
  permutations = list(itertools.permutations(input))
  output = []
  for n in range(0,len(permutations)):
    if (permutations[n][4] != "green" and permutations[n][0] != "white" and permutations[n][1] == "blue"):
      output.append(list(permutations[n]))
  print "Remaining items in colors " + str(len(output))
  return output

#  8 The man living in the house right in the center drinks milk.
#  5 The green house owner drinks coffee.
# Inputting (drinks) makes sure milk is always in the middle. Green cannot be last so coffee cannot be last.
def drinkConditions(input):
  permutations = list(itertools.permutations(input))
  output = []
  for n in range(0,len(permutations)):
    if (permutations[n][2] == "milk" and permutations[n][4] != "coffee"):
      output.append(list(permutations[n]))
  print "Remaining items in drinks " + str(len(output))
  return output

# Method validate(composition) checks any remaining compositions against the remaining business rules.
def validate(composition):

# Checking each individual house until any deviation is found in case of which False is returned.
  for x in range(0,5):

#  1 The Brit lives in a red house.
    if (composition[1][x] == "UK" and composition[2][x] != "red"):
      return False

#  2 The Swede keeps dogs as pets.
    if (composition[1][x] == "S" and composition[5][x] != "dogs"):
      return False

#  3 De Deen drinkt thee.
    if (composition[1][x] == "DK" and composition[3][x] != "tea"):
      return False

#  5 The green house owner drinks coffee.
    if (composition[2][x] == "green" and composition[3][x] != "coffee"):
      return False

#  6 The person who smokes Pall Mall rears birds.
    if (composition[4][x] == "Pall Mall" and composition[5][x] != "birds"):
      return False

#  7 De eigenaar van het gele huis rookt Dunhill.
    if (composition[2][x] == "yellow" and composition[4][x] != "Dunhill"):
      return False

# 12 The owner who smokes Blue Master drinks beer.
    if (composition[4][x] == "Blue Master" and composition[3][x] != "beer"):
      return False

# 13 The German smokes Prince.
    if (composition[1][x] == "DE" and composition[4][x] != "Prince"):
      return False

# And checking the composition to see if the neighbour rules apply.

#  4 The green house is on the left of the white, next to it.
  if((findIndex(2,"green",composition) - findIndex(2,"white",composition)) != -1):
    return False

# 10 The man who smokes blend lives next to the one who keeps cats.
  if(abs(findIndex(4,"blend",composition) - findIndex(5,"cats",composition)) != 1):
    return False

# 11 The man who keeps horses lives next to the man who smokes Dunhill.
  if(abs(findIndex(5,"horse",composition) - findIndex(4,"Dunhill",composition)) != 1):
    return False

# 15 The man who smokes blend has a neighour who drinks water.
  if(abs((findIndex(4,"blend",composition) - findIndex(3,"water",composition))) != 1):
    return False

# If no reason to return False, composition must be the right one, so True is returned.
  return True

# Method findIndex(position,value,composition) looks up the house number of a certain value in a certain composition in order to help do the neighbour business rules checks.
def findIndex(position,value,composition):
  for x in range(0,5):
    if (composition[position][x] == value):
      return composition[0][x]

# The actual program.
def init():

  print "Pre-selecting..."

# Tweaking the lists of permutations where possible.
  allCountries = countryFirst(countries,"NO")
  allColors = colorConditions(colors)
  allDrinks = drinkConditions(drinks)
  allSmokes = list(itertools.permutations(smokes))
  allPets = list(itertools.permutations(pets))

  tries = 0
  toTry = len(allCountries)*len(allColors)*len(allDrinks)*len(allSmokes)*len(allPets)

  print "Starting validating " + str(toTry) + " options..."

# Running validations for each composition.
  for a in range(0,len(allCountries)):
    print "Validating, progress...(" + str(round((float(tries) / float(toTry) * 100),1)) + "%)"
    for b in range(0,len(allColors)):
      for c in range(0,len(allDrinks)):
        for d in range(0,len(allSmokes)):
          for e in range(0,len(allPets)):
            composition = [(1,2,3,4,5),allCountries[a],allColors[b],allDrinks[c],allSmokes[d],allPets[e]]
            tries += 1
            if(validate(composition) == True):
              winner = composition

# Prints the winner in a kind-of-neat format.
  for m in range(0,5):
    for n in range(0,6):
      print winner[n][m]

# Runs the actual program after all methods are defined.
init()