Given a set of N numbers (operands) and a M binary operators, finds all N-ary expressions evaluating to specified target value.

mathquiz.py

"""Given a set of N numbers (operands) and a M binary operators,
   finds all N-ary expressions evaluating to specified target value.
   See this fragment of the Countdown TV show for a sample problem:
   http://www.youtube.com/watch?v=EO472qM6M9g"""

import itertools
import sys

class Leaf:
    def __init__(self):
        self.is_leaf = True

    def visit_inorder(self, func):
        func(self)

class Node:
    def __init__(self, left, right):
        self.is_leaf = False
        self.left = left
        self.right = right

    def visit_inorder(self, func):
        self.left.visit_inorder(func)
        func(self)
        self.right.visit_inorder(func)

def binary_trees(num_nodes):
    if num_nodes == 0:
        yield Leaf()
    for i in xrange(num_nodes):
        for left_subtree in binary_trees(i):
            for right_subtree in binary_trees(num_nodes - i - 1):
                yield Node(left_subtree, right_subtree)

def assign_node_operators(tree, operator_list):
    op_iter = iter(operator_list)
    def assign_node_operator(node):
        if not node.is_leaf:
            node.operator = next(op_iter)
    tree.visit_inorder(assign_node_operator)

def assign_leaf_values(tree, value_list):
    value_iter = iter(value_list)
    def assign_leaf_value(node):
        if node.is_leaf:
            node.value = next(value_iter)
    tree.visit_inorder(assign_leaf_value)

def evaluate_expr_tree(tree):
    def evaluate_node(node):
        if not node.is_leaf:
            return node.operator(
                evaluate_node(node.left),
                evaluate_node(node.right))
        else:
            return node.value
    return evaluate_node(tree)

def print_tree_topology(tree):
    def print_node(node):
        if not node.is_leaf:
            sys.stdout.write('(')
            print_node(node.left)
            print_node(node.right)
            sys.stdout.write(')')
        else:
            sys.stdout.write('x')
    print_node(tree)
    print

def print_expr_tree(tree):
    def print_node(node):
        if not node.is_leaf:
            sys.stdout.write('(')
            print_node(node.left)
            sys.stdout.write(node.operator.symbol)
            print_node(node.right)
            sys.stdout.write(')')
        else:
            sys.stdout.write(str(node.value))
    print_node(tree)
    print

class BinaryOperator:
    def __init__(self, symbol, func):
        self.symbol = symbol
        self.func = func

    def __call__(self, a, b):
        return self.func(a, b)

def solve_math_quiz(target, numbers, operators):
    num_ops = len(numbers) - 1
    for tree in binary_trees(num_ops):
        print 'Topology: ',
        print_tree_topology(tree)
        for op_comb in itertools.product(operators, repeat=num_ops):
            assign_node_operators(tree, op_comb)
            for num_perm in itertools.permutations(numbers):
                assign_leaf_values(tree, num_perm)
                try:
                    result = evaluate_expr_tree(tree)
                    if result == target:
                        print_expr_tree(tree)
                except:
                    # catch divisions by zero, etc.
                    pass


if __name__ == '__main__':
    numbers = [ 25, 50, 75, 100, 3, 6 ]
    operators = [
        BinaryOperator('+', lambda a, b: a + b),
        BinaryOperator('-', lambda a, b: a - b),
        BinaryOperator('*', lambda a, b: a * b),
        BinaryOperator('/', lambda a, b: a / b),
        #BinaryOperator('%', lambda(a, b): a % b),
    ]

    solve_math_quiz(952, numbers, operators)