Matchstick equation solver

matchstick_equations.py

from itertools import product

SYMBOLS = {
    '0': (1, 1, 1, 1, 1, 1, 0, 0, 0, 0),
    '1': (0, 0, 0, 1, 1, 0, 0, 0, 0, 0),
    '2': (1, 0, 1, 1, 0, 1, 1, 0, 0, 0),
    '3': (0, 0, 1, 1, 1, 1, 1, 0, 0, 0),
    '4': (0, 1, 0, 1, 1, 0, 1, 0, 0, 0),
    '5': (0, 1, 1, 0, 1, 1, 1, 0, 0, 0),
    '6': (1, 1, 1, 0, 1, 1, 1, 0, 0, 0),
    '7': (0, 0, 1, 1, 1, 0, 0, 0, 0, 0),
    '8': (1, 1, 1, 1, 1, 1, 1, 0, 0, 0),
    '9': (0, 1, 1, 1, 1, 1, 1, 0, 0, 0),
    '+': (0, 0, 0, 0, 0, 0, 1, 1, 0, 0),
    '-': (0, 0, 0, 0, 0, 0, 1, 0, 0, 0),
    '=': (0, 0, 0, 0, 0, 0, 0, 0, 1, 1)
}

SYMBOLS_REVERSE = {value: key for key, value in SYMBOLS.items()}


def memoize(f):
    class MemoDict(dict):
        def __init__(self, f):
            self.f = f

        def __call__(self, *args):
            return self[args]

        def __missing__(self, key):
            res = self[key] = self.f(*key)
            return res

    return MemoDict(f)


@memoize
def remove_or_add_stick(c, b):
    a = SYMBOLS[c]
    neighbors = set()
    for i in range(len(a)):
        if a[i] == b:
            continue
        copy = list(a)
        copy[i] = b
        copy = tuple(copy)
        if copy in SYMBOLS_REVERSE:
            neighbors.add(copy)
    return neighbors


@memoize
def move_stick(c):
    a = SYMBOLS[c]
    ones = (i for i, e in enumerate(a) if e == 1)
    zeros = (i for i, e in enumerate(a) if e == 0)
    neighbors = set()
    for i, j in product(ones, zeros):
        perm = list(a)
        perm[i], perm[j] = perm[j], perm[i]
        perm = tuple(perm)
        if perm in SYMBOLS_REVERSE:
            neighbors.add(SYMBOLS_REVERSE[perm])
    return neighbors


def solve(equation_str):
    """
    Solves matchstick equation puzzles where a single matchstick has to be moved. Stops at the first solution.
    """
    for first, c1 in enumerate(equation_str):
        if c1 == ' ':
            continue
        neighbors = move_stick(c1)
        for neighbor in neighbors:
            equation_copy = list(equation_str)
            equation_copy[first] = neighbor
            if eval(''.join(equation_copy).replace('=', '==')):
                return ''.join(equation_copy)
        rem_neighbors = remove_or_add_stick(c1, 0)
        for rem_neighbor in rem_neighbors:
            first_symbol = SYMBOLS_REVERSE[rem_neighbor]
            for second, c2 in enumerate(equation_str):
                if c2 == ' ':
                    continue
                add_neighbors = remove_or_add_stick(c2, 1)
                for add_neighbor in add_neighbors:
                    second_symbol = SYMBOLS_REVERSE[add_neighbor]
                    equation_copy = list(equation_str)
                    equation_copy[first] = first_symbol
                    equation_copy[second] = second_symbol
                    if eval(''.join(equation_copy).replace('=', '==')):
                        return ''.join(equation_copy)
    return 'NO SOLUTION!'


if __name__ == '__main__':
    eqs = [
        '185 + 15 = 270',
        '8 - 2 = 1',
        '2 + 2 = 5',
        '1 + 1 = 20'
    ]
    for eq in eqs:
        print(solve(eq))