Who owns the zebra?

I was prompted to add a Python code solution to the classic Zebra puzzle after working on two logic puzzles on the web that have a similar form: Sunday Times Teaser 2606 and New Scientist Enigma 1430.

The code to solve the Zebra Puzzle adds one category (house colour, nationality, drink, smoke, pet) at a time, filtering as much as possible at each step, in order to minimise the number of combinations that need to be checked:



from itertools import permutations as perm

def dict_perms( category_list ):
  for p in perm(range(1,6)): yield {a:b for (a,b) in zip(category_list,p)}

def key(d, v): # find a key in dictionary d with value v
  return [name for name, val in d.items() if val == v][0]

colours = [ colour for colour in dict_perms(['Red','Green','Indigo','Yellow','Blue'])
  if colour['Green'] == colour['Indigo']+1 ]

colour_nation = [ (colour,nation) for colour in colours
  for nation in dict_perms(['Englishman','Spaniard','Ukranian','Norweigan','Japanese'])
  if  nation['Norweigan']==1
  and colour['Red'] == nation['Englishman'] 
  and abs(nation['Norweigan'] - colour['Blue'])==1 ]

colour_nation_drink = [(colour, nation, drink) for colour,nation in colour_nation
  for drink in dict_perms(['Coffee','Milk','Orange Juice','Water','Tea'])
  if  drink['Coffee'] == colour['Green']
  and nation['Ukranian'] == drink['Tea']
  and drink['Milk']==3 ]

colour_nation_drink_smoke = [(colour, nation, drink, smoke)
  for colour, nation, drink in colour_nation_drink
  for smoke in dict_perms(['Old Gold','Kools','Chesterfield','Lucky Strike','Parliament']) 
  if  smoke['Kools'] == colour['Yellow']
  and smoke['Lucky Strike'] == drink['Orange Juice']
  and nation['Japanese'] == smoke['Parliament']     ]

solutions = [(colour, nation, drink, smoke, pet)
  for colour, nation, drink, smoke in colour_nation_drink_smoke
  for pet in dict_perms(['Dog','Snails','Fox','Horse','Zebra']) 
  if  smoke['Old Gold'] == pet['Snails']
  and abs(smoke['Chesterfield'] - pet['Fox'])==1
  and abs(smoke['Kools'] - pet['Horse'])==1
  and nation['Spaniard'] == pet['Dog']   ]

if len(solutions)>1:
  print "multiple solutions"
else:
  (colour,nation,drink,smoke,pet) = solutions[0]
  for pos in range(1,6):
    print pos, ":", key(colour,pos), ",",key(nation,pos), ",",key(drink,pos),",",key(smoke,pos), ",",key(pet,pos)