538 Riddler Express: Solving a festive cryptarithm!

cryptarithm.py

# Solution for https://fivethirtyeight.com/features/can-you-outwit-the-tax-collector/
from constraint import *


def unique(collection):
    """Gets a sorted list of unique elements from a collecion."""
    return sorted(list(set(collection)))


def word(*digits):
    """Converts a list of digits into a base 10 integer."""
    return sum([int(digit) * 10 ** i for i, digit in enumerate(reversed(digits))])


def format_solution(solution):
    """Formats a solution ordered by variable names."""
    keys = sorted(solution.keys())
    return ", ".join([f"{k}={solution[k]}" for k in keys])


def dynamic_constraint(words):
    """Creates a dynamic CSP constraint from a list of words."""
    letters = unique("".join(words))
    word_exprs = [f"word({','.join(list(word))})" for word in words]
    params = ",".join(letters)
    lhs = " + ".join(word_exprs[:-1])
    rhs = word_exprs[-1]
    constraint_fn = eval(f"lambda {params}: {lhs} == {rhs}")
    return constraint_fn, letters


def solve_cryptarithm(words):
    """Assume that the last word in the list is equal to the sum of the others."""
    constraint_fn, letters = dynamic_constraint(words)

    # No leading zeros allowed
    initial_letters = unique([word[0] for word in words if len(word) > 1])
    other_letters = [letter for letter in letters if letter not in initial_letters]

    problem = Problem()
    problem.addVariables(initial_letters, list(range(1, 10)))
    problem.addVariables(other_letters, list(range(10)))
    problem.addConstraint(AllDifferentConstraint())
    problem.addConstraint(constraint_fn, letters)
    solutions = problem.getSolutions()
    return solutions


# solutions = solve_cryptarithm(["TO", "GO", "OUT"])
# solutions = solve_cryptarithm(["SEND", "MORE", "MONEY"])
solutions = solve_cryptarithm(["HAPPY", "HOLIDAYS", "HOHOHOHO"])

if len(solutions) == 0:
    print("No solutions found")
for solution in solutions:
    print(format_solution(solution))
    
# Two solutions:
# A=4, D=3, H=8, I=2, L=7, O=0, P=6, S=9, Y=1
# A=8, D=7, H=6, I=4, L=5, O=1, P=3, S=9, Y=2