fourfours.py

import operator
import sys
import argparse
from collections import namedtuple

# Parse command line arguments
parser = argparse.ArgumentParser(description='Generate solutions to the four fours problem')
parser.add_argument('-d', '--display', choices=['number', 'expression', 'both'], default='both', help='What to display for each result')
parser.add_argument('-c', '--count', type=int, default=4, help='The number of fours to use')
parser.add_argument('-n', '--number', type=int, default=4, help='The number to use in expressions')
parser.add_argument('-p', '--positive', action='store_true', help='Only output positive numbers')
parser.add_argument('-t', '--use-concat', action='store_true', help='Allow concatenating numbers (i.e. 4 and 4 to 44)')
parser.add_argument('-m', '--max-exponent', type=int, default=4096, help='The maximum number to allow as an exponent')
args = parser.parse_args()

# Initialize operators to use
Operator = namedtuple('Operator', ['function', 'symbol', 'precedence', 'associativity', 'commutative'])

add = Operator(operator.add, '+', 3, 'left', True)
sub = Operator(operator.sub, '-', 3, 'left', False)
mul = Operator(operator.mul, '*', 2, 'left', True)
div = Operator(operator.floordiv, '/', 2, 'left', False)
pow = Operator(operator.pow, '^', 1, 'right', False)
concat = Operator(lambda a, b: 10*a+b, '', 0, 'left', True)

operators = [add, sub, mul, div, pow]

if args.use_concat:
    operators.append(concat)

# Class representing an operator in the expression tree
class Node:
    def __init__(self, op, left, right):
        self.op = op
        self.left, self.right = left, right

    # Evaluate the expression
    def evaluate(self):
        if type(self.left) == Node:
            left = self.left.evaluate()
        else:
            left = self.left

        if type(self.right) == Node:
            right = self.right.evaluate()
        else:
            right = self.right

        # Check arguments to '^' to avoid very large or fractional results
        if self.op == pow and right > args.max_exponent or right < 0:
            raise ValueError
        # Check arguments to '/' to avoid divinding by 0 or fractional results
        if self.op == div and (right == 0 or left % right != 0):
            raise ValueError
        return self.op.function(left, right)

    # Return the expression as a string
    def __str__(self):
        left = str(self.left)
        right = str(self.right)
        # Parenthesize argumements based on precedence, associativity, and commutativity
        if (type(self.left) == Node and
                (self.left.op.precedence > self.op.precedence or
                    (self.left.op.precedence == self.op.precedence and
                        self.left.op.associativity == 'right' and not self.op.commutative))):
            left = '(' + left +')'
        if (type(self.right) == Node and
                (self.right.op.precedence > self.op.precedence or
                    (self.right.op.precedence == self.op.precedence and
                        self.right.op.associativity == 'left' and not self.op.commutative))):
            right = '(' + right +')'
        return left + self.op.symbol + right

# Generator to produce all expression trees.
# Modified from http://nedbatchelder.com/blog/200802/enumerating_trees.html
def trees(stack, fours):
    if len(stack) == 1 and fours == 0:
        yield stack[0]
    if len(stack) > 1:
        for op in operators:
            left = stack[-2]
            right = stack[-1]
            if op.commutative and type(right) == Node and right.op == op:
                continue
            if (op == concat and ((type(right) == Node and right.op != concat) or
                    (type(left) == Node and left.op != concat))):
                continue
            new_node = Node(op, left, right)
            new_stack = stack[:-2]
            new_stack.append(new_node)
            for t in trees(new_stack, fours):
                yield t
    if fours >= 0:
        stack.append(args.number)
        for t in trees(stack, fours - 1):
            yield t

# Dictionary to store results in
results = {}

# Loop through expressions, adding new ones to results
for tree in trees([], args.count):
    try:
        result = tree.evaluate()
    except (ValueError):
        continue
    if result in results:
        continue
    results[result] = str(tree)

# For integer to string conversions, first try to use gmpy2 for performance, then fallback to the standard library
# Using gmpy2 is highly recommended if --count is increased
try:
    from gmpy2 import digits as int_to_str
except ImportError:
    int_to_str = str

# Loop through results to print. Key function puts positives first then orders by magnitude.
for number in sorted(results.keys(), key=lambda n:(n<0, abs(n))):
    if number < 0 and args.positive:
        continue
    output = ''
    if args.display in ('number', 'both'):
        output += int_to_str(number)
    if args.display  == 'both':
        output += ': '
    if args.display in ('expression', 'both'):
        output += results[number]
    print(output)