solve_mastermind.py

from collections import defaultdict
from itertools import permutations
from pprint import pprint
# this is going to calculate the output of the game
# the number of colors guessed correctly and the number of 
# colors placed correctly
def calculate_score(correct, guess):
    correct_colors = sum(1 for color in guess if color in correct)
    correct_place = sum(1 for (a, b) in zip(correct, guess) if a == b)
    return (correct_colors - correct_place, correct_place)

# print(calculate_score((0,1,2,3), (1,0,2,5)))
colors = list(range(8))
possible = list(permutations(colors, 4))

def calculate_distribution(remaining, guess):
    distribution = defaultdict(list)
    for correct in remaining:
        score = calculate_score(correct, guess)
        distribution[score].append(correct)
    return distribution

def find_best_guess(remaining):
    best_guess = None
    best_score = len(possible) + 1
    for guess in possible:
        distribution = calculate_distribution(remaining, guess)
        dist_score = max(len(solutions) for solutions in distribution.values())
        if dist_score < best_score:
            best_score = dist_score
            best_guess = guess
    # print('best score:', best_score)
    return best_guess

starting_guess = (0,1,2,3)

def make_decision_tree(remaining, guess):
    tree = dict()
    distribution = calculate_distribution(remaining, guess)
    for (score, next_remaining) in distribution.items():
        if len(next_remaining) <= 1:
            tree[score] = {
                'solution': next_remaining[0]
            }
        else:
            best_guess = find_best_guess(next_remaining)
            tree[score] = {
                'best_guess': best_guess,
                'solution_space': len(next_remaining),
                'children': make_decision_tree(next_remaining, best_guess)
            }
    return tree

tree = make_decision_tree(possible, starting_guess)

def print_tree(tree, depth=0):
    padding = ' ' * (depth * 4)
    for (score, child) in tree.items():
        if 'solution' in child:
            print (f'{padding}score={score} solution={child["solution"]}')
        if 'best_guess' in child:
            print(f'{padding}score={score} best_guess={child["best_guess"]} solution_space={child["solution_space"]}')
            print_tree(child['children'], depth + 1)
print_tree(tree)