3 solver implementations for send+more=money problem

send_more_money.py

from typing import Optional, Tuple, Set, List, Dict
import copy
import time
from argparse import ArgumentParser
import random
import json
import numpy as np
import os
import uuid
from types import SimpleNamespace


def check_solution(d: dict) -> bool:
    """
    Check that the equation works
    We are looking for the one solution where none of the leading digits are zeros
    """
    unique_vals = set([])
    for v in d.values():
        unique_vals.add(v)
    if len(unique_vals) < 8:
        return False

    if d["m"] == 0 or d["s"] == 0:
        return False

    return (d["s"] * 1000 + d["e"] * 100 + d["n"] * 10 + d["d"]) + (
        d["m"] * 1000 + d["o"] * 100 + d["r"] * 10 + d["e"]
    ) == (d["m"] * 10_000 + d["o"] * 1000 + d["n"] * 100 + d["e"] * 10 + d["y"])


def dfs_uniq(
    assigned: dict, unassigned_letters: List[str], unassigned_numbers: Set[int]
) -> Tuple[int, Optional[dict]]:
    """Search mechanism for brute_force_uniq technique. Use DFS for good memory performance."""
    if unassigned_letters == []:
        if check_solution(assigned):
            return 1, assigned
        else:
            return 1, None

    num_checked = 0
    var = unassigned_letters[0]  # highly dependent on initial variable ordering
    assign_prime = copy.copy(assigned)
    for val in unassigned_numbers:
        assign_prime[var] = val
        unassigned_n_prime = unassigned_numbers - set([val])
        n, solution = dfs_uniq(
            assigned=assign_prime,
            unassigned_letters=unassigned_letters[1:],
            unassigned_numbers=unassigned_n_prime,
        )
        num_checked += n
        if solution:
            return num_checked, solution
    return num_checked, None  # solution not found


def print_solution(solution: Optional[dict]) -> None:
    if solution:
        print("found solution:")
        d = solution
        top_row = " %d%d%d%d" % (d["s"], d["e"], d["n"], d["d"])
        print(top_row)
        next_row = "+%d%d%d%d" % (d["m"], d["o"], d["r"], d["e"])
        print(next_row)
        print("-----")
        last_row = "%d%d%d%d%d" % (d["m"], d["o"], d["n"], d["e"], d["y"])
        print(last_row)
    else:
        print("no solution found")


def print_stats(num_checked: int, time_taken: float) -> None:
    print(f"# solutions tried: {num_checked:,}")
    print(f"took {time_taken:.2f} seconds")


def find_soln_brute_force_uniq(use_m_1: bool = False, letter_order: str = "sorted"):
    print("searching using technique 'uniq_brute_force'...")
    unassigned_letters, assigned = init(use_m_1, letter_order)
    unassigned_numbers = set([i for i in range(0, 10)])

    if use_m_1:
        unassigned_numbers.remove(1)

    start = time.time()
    num_checked, soln = dfs_uniq(assigned, unassigned_letters, unassigned_numbers)
    end = time.time()

    print_solution(soln)
    print_stats(num_checked, end - start)
    return num_checked, end - start


def dfs_brute(assigned: dict, unassigned: List[str]) -> Tuple[int, Optional[dict]]:
    if unassigned == []:
        if check_solution(assigned):
            return 1, assigned
        else:
            return 1, None

    num_checked = 0
    var = unassigned[0]
    assign_prime = copy.copy(assigned)

    for val in range(10):
        assign_prime[var] = val
        n, soln = dfs_brute(assigned=assign_prime, unassigned=unassigned[1:],)
        num_checked += n
        if soln:
            return num_checked, soln
    return num_checked, None


def find_soln_brute_force(use_m_1: bool = False, letter_order: str = "sorted"):
    print("searching using technique 'brute_force'...")
    letters, assigned = init(use_m_1, letter_order)

    start = time.time()
    num_checked, soln = dfs_brute(assigned=assigned, unassigned=letters,)
    end = time.time()

    print_solution(soln)
    print_stats(num_checked, end - start)
    return num_checked, end - start


class Constraint:
    def can_check(self, assignment: dict) -> bool:
        raise NotImplementedError

    def check(self, assignment: dict) -> bool:
        raise NotImplementedError


class Mod10Constraint(Constraint):
    def __init__(self, a, b, c):
        self.a = a
        self.b = b
        self.c = c

    def can_check(self, assignment: dict) -> bool:
        return self.a in assignment and self.b in assignment and self.c in assignment


class StrictMod10Constraint(Mod10Constraint):
    """the constraint s.t. (a + b) % 10 == c"""

    def check(self, assignment: dict) -> bool:
        return (assignment[self.a] + assignment[self.b]) % 10 == assignment[self.c]


class LooseMod10Constraint(Mod10Constraint):
    """the constraint s.t. (a + b) % 10 == c or (a + b + 1) % 10 == c"""

    def check(self, assignment: dict) -> bool:
        return (assignment[self.a] + assignment[self.b]) % 10 == assignment[self.c] or (
            assignment[self.a] + assignment[self.b] + 1
        ) % 10 == assignment[self.c]


class BinaryConstraint(Constraint):
    def __init__(self, letter):
        self.letter = letter

    def can_check(self, assignment: dict):
        return self.letter in assignment

    def check(self, assignment: dict):
        return assignment[self.letter] in [0, 1]


class NonZeroLeading(Constraint):
    def __init__(self, letter):
        self.letter = letter

    def can_check(self, assignment: dict):
        return self.letter in assignment

    def check(self, assignment: dict):
        return assignment[self.letter] != 0


def dfs_with_constraints(
    assigned: dict,
    unassigned_letters: List[str],
    unassigned_numbers: Set[int],
    constraints: dict,
) -> Tuple[int, Optional[dict]]:
    """
    Run depth-first search with FC (forward-checking)
    Rely on the constraints
    """

    if unassigned_letters == []:
        # all variables are assigned
        if check_solution(assigned):
            return 1, assigned
        else:
            return 1, None
    else:
        num_checked = 0

        # pick the variable which has the most constraints
        unassigned_letters.sort(key=lambda letter: -1 * len(constraints[letter]))

        # pick an unassigned variable
        var = unassigned_letters[0]

        assigned_prime = copy.copy(assigned)
        for val in unassigned_numbers:
            assigned_prime[var] = val

            is_dwo = False
            for constraint in constraints[var]:
                if constraint.can_check(assigned_prime):
                    if not constraint.check(assigned_prime):
                        # domain wipeout
                        is_dwo = True
                        break
            if is_dwo:
                continue

            # otherwise (NO ELSE)
            unassigned_n_prime = unassigned_numbers - set([val])
            n, soln = dfs_with_constraints(
                assigned=assigned_prime,
                unassigned_letters=unassigned_letters[1:],
                unassigned_numbers=unassigned_n_prime,
                constraints=constraints,
            )
            num_checked += n
            if soln:
                return num_checked, soln

        return num_checked, None


def find_soln_constraints(use_m_1: bool = False, letter_order: str = "sorted"):
    print("searching using technique 'constraints'...")

    letters, starting_assignment = init(use_m_1, letter_order)
    unassigned_numbers = set([i for i in range(10)])
    if use_m_1:
        unassigned_numbers.remove(1)
    constraint_s_m_o = LooseMod10Constraint("s", "m", "o")
    constraint_e_o_n = LooseMod10Constraint("e", "o", "n")
    constraint_n_r_e = LooseMod10Constraint("n", "r", "e")
    constraint_d_e_y = StrictMod10Constraint("d", "e", "y")
    bin_constraint_m = BinaryConstraint("m")
    non_zero_leading_s = NonZeroLeading("s")
    non_zero_leading_m = NonZeroLeading("m")
    constraints = {
        "s": [constraint_s_m_o, non_zero_leading_s],
        "e": [constraint_n_r_e, constraint_d_e_y, constraint_e_o_n],
        "n": [constraint_e_o_n, constraint_n_r_e],
        "d": [constraint_d_e_y],
        "m": [constraint_s_m_o, bin_constraint_m, non_zero_leading_m],
        "o": [constraint_s_m_o, constraint_e_o_n],
        "r": [constraint_n_r_e],
        "y": [constraint_d_e_y],
    }

    start = time.time()
    num_checked, soln = dfs_with_constraints(
        assigned=starting_assignment,
        unassigned_letters=letters,
        unassigned_numbers=unassigned_numbers,
        constraints=constraints,
    )
    end = time.time()

    print_solution(soln)
    print_stats(num_checked, end - start)
    return num_checked, end - start


def init(
    use_m_1: bool = False, letter_order: str = "sorted"
) -> Tuple[List[str], Dict[str, int]]:
    """Simplify handling arguments"""
    letters = list(set([c for c in "sendmoremoney"]))
    if letter_order == "pathological":
        print("[init] using pathological letter ordering")
        # to trigger the worst case I think we need to use this assignment:
        letters = ["s", "r", "d", "n", "e", "y", "m", "o"]
    elif letter_order == "random":
        print("[init] using random letter ordering")
        random.shuffle(letters)
    elif letter_order == "sorted":
        print("[init] using sorted letter ordering")
        letters.sort()

    assigned = {}
    if use_m_1:
        print("[init] using m=1 trick")
        assigned["m"] = 1
        letters.remove("m")

    return (letters, assigned)


def main(args):
    if args.solver == "brute_force":
        solver = find_soln_brute_force
    elif args.solver == "brute_force_uniq":
        solver = find_soln_brute_force_uniq
    elif args.solver == "fc":
        solver = find_soln_constraints
    else:
        raise Exception("unknown solver: %s" % args.solver)

    times = []
    states = []
    for i in range(args.num_experiments):
        num_states, time_taken = solver(
            use_m_1=args.use_m_1,
            letter_order=args.letter_order,
        )
        states.append(num_states)
        times.append(time_taken)

    agg = {}
    if len(times) > 1:
        print("-" * 40)
        print(f"ran search {args.num_experiments} times")
        agg["mean_time"] = np.mean(times)
        agg["stdev_time"] = np.std(times)
        print(f"mean time: {np.mean(times):.2f} ± {np.std(times):.2f}")
        agg["mean_states"] = np.mean(states)
        agg["stdev_states"] = np.std(states)
        print(f"mean states: {np.mean(states):,.2f} ± {np.std(states):,.2f}")
        print("-" * 40)

    return times, states, agg


def try_a_bunch_of_stuff():
    """Run all techniques and variable orders in a batch to generate performance stats"""
    try:
        os.mkdir("data")
    except FileExistsError:
        pass

    run_id = uuid.uuid4().hex
    print(f"starting Run ID {run_id}")
    run_order = 1

    for solver in ["brute_force", "brute_force_uniq", "fc"]:
        fname = f"data/run-{run_id}-{solver}.json"
        runs = []

        for use_m_1 in [False, True]:
            for letter_order in ["pathological", "random"]:
                run_params = {
                    "letter_order": letter_order,
                    "use_m_1": use_m_1,
                    "solver": solver,
                    "run_id": run_id,
                    "run_order": run_order,
                    "num_experiments": 5,
                }
                args = SimpleNamespace(name="foo", **run_params)

                times, states, agg = main(args)
                runs.append(
                    {
                        "run_params": run_params,
                        "results": {"times": times, "states": states, "agg": agg,},
                    }
                )
                run_order += 1

        with open(fname, "w") as fp:
            json.dump(runs, fp, indent=4)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "-s",
        "--solver",
        required=True,
        choices=["brute_force", "brute_force_uniq", "fc"],
    )
    parser.add_argument(
        "--use-m-1",
        default=False,
        action="store_true",
        help="assign m=1 to speed up computation",
    )
    parser.add_argument(
        "-o",
        "--letter-order",
        default="sorted",
        choices=["random", "pathological", "sorted"],
        help="set order of letters to explore",
    )
    parser.add_argument("-n", "--num-experiments", type=int, default=1)
    args = parser.parse_args()

    main(args)

    # try_a_bunch_of_stuff()