Das deutsche Python-Forum

Code: Alles auswählen

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""This module solves Einstein's riddle.

Requires Gustavo Niemeyer's `constraint` module.
"""
import time
import pickle
from constraint import Problem, AllDifferentConstraint

__author__ = "Marc 'BlackJack' Rintsch"
__version__ = '0.1.0'
__date__ = '$Date: 2005-08-22 13:20:15 +0200 (Mon, 22 Aug 2005) $'
__revision__ = '$Rev: 747 $'


def varnames(name):
    """Creates a list of variable names by adding a number suffix from 1 to 5
    to the `name`.
    """
    return ['%s_%d' % (name, i + 1) for i in xrange(5)]


def relates(value_a, value_b):
    """Returns a constraint function that tests if two variables from different
    domains but for the same house contain `value_a` and `value_b`.
    """
    def relates_constraint(*args):
        """Takes 10 arguments, 5 values for the first domain and 5 for the
        second.  The values are ordered by house number.
        """
        for item_a, item_b in zip(args[:5], args[5:]):
            if item_a == value_a and item_b == value_b:
                return True
        return False
    return relates_constraint


def neighbour(value_a, value_b):
    """Returns a constraint function that checks if the house that has
    `value_a` is next to a house that has `value_b`.
    """
    def neighbour_constraint(*args):
        """Takes 10 arguments, 5 values for the first domain and 5 for the
        second.  The values are ordered by house number.
        """
        a_values = args[:5]
        b_values = args[5:]
        for item_a, item_b in (zip(a_values, b_values[1:])
                               + zip(a_values[1:], b_values)):
            if item_a == value_a and item_b == value_b:
                return True
        return False
    return neighbour_constraint


def green_left_of_white(*args):
    """Tests if the green house has a lower house number than the white
    house.
    """
    house = dict((color, house_nr) for house_nr, color in enumerate(args))
    return house['green'] < house['white']


def main():
    """Main function."""
    # 
    # Variables.
    # 
    nationality_vars = varnames('nationality')
    color_vars = varnames('color')
    drink_vars = varnames('drink')
    cigarette_vars = varnames('cigarette')
    pet_vars = varnames('pet')

    # 
    # Domains.
    # 
    nationalities = ('british', 'danish', 'german', 'norwegian', 'swedish')
    colors = ('blue', 'green', 'red', 'white', 'yellow')
    drinks = ('beer', 'coffee', 'milk', 'tea', 'water')
    cigarettes = ('Dunhill', 'Pall Mall', 'Malboro', 'Rothmanns', 'Winfield')
    pets = ('bird', 'cat', 'dog', 'fish', 'horse')
    
    # 
    # Houston, we have a problem.
    # 
    problem = Problem()
    
    # 
    # Add the variables and their domains to the problem.
    # 
    for variables, domain in ((nationality_vars, nationalities),
                              (color_vars, colors),
                              (drink_vars, drinks),
                              (cigarette_vars, cigarettes),
                              (pet_vars, pets)):
        problem.addVariables(variables, domain)
    
    problem.addConstraint(AllDifferentConstraint())
    
    # 
    # The hints of the riddle as tuples of constraint function and sequence
    # of involved variables.
    # 
    constraints = ((relates('british', 'red'), nationality_vars + color_vars),
                   (relates('swedish', 'dog'), nationality_vars + pet_vars),
                   (relates('danish', 'tea'), nationality_vars + drink_vars),
                   (green_left_of_white, color_vars),
                   (relates('green', 'coffee'), color_vars + drink_vars),
                   (relates('Pall Mall', 'bird'), cigarette_vars + pet_vars),
                   (lambda a: a == 'milk', ['drink_3']),
                   (relates('yellow', 'Dunhill'), color_vars + cigarette_vars),
                   (lambda a: a == 'norwegian', ['nationality_1']),
                   (neighbour('Malboro', 'cat'), cigarette_vars + pet_vars),
                   (neighbour('horse', 'Dunhill'), pet_vars + cigarette_vars),
                   (relates('Winfield', 'beer'), cigarette_vars + drink_vars),
                   (neighbour('norwegian', 'blue'),
                    nationality_vars + color_vars),
                   (relates('german', 'Rothmanns'),
                    nationality_vars + cigarette_vars),
                   (neighbour('Malboro', 'water'), cigarette_vars + drink_vars))
    
    for constraint_func, variables in constraints:
        problem.addConstraint(constraint_func, variables)
    
    # 
    # Solve the problem and save the solutions as pickle.
    # 
    start_time = time.time()
    
    print 'Searching solutions...'
    solutions = list()
    for solution in problem.getSolutionIter():
        print solution
        solutions.append(solution)
    
    elapsed_time = time.time() - start_time
    print 'Found %d solutions in %d seconds.' % (len(solutions), elapsed_time)
    
    solution_file = open('solutions.dat', 'wb')
    pickle.dump(dict(solutions=solutions), solution_file)
    solution_file.close()


if __name__ == '__main__':
    main()