Solving an arithmetic puzzle with Haskell

[EDIT: since this post still seems to get a good deal of traffic, I should note that (as you can see if you read the comments) the code I gave here is not quite correct. Still, it’s interesting enough that I’ll leave it up.]

JD2718 posted a puzzle the other day: the idea is to count how many possible results you can get by combining the numbers 4,3,2,1 (in that order) with the four arithmetic operators and parentheses. Naturally I decided to write some Haskell code to solve this one.

First, instead of thinking in terms of possibly using parentheses, I just generate all possible postfix expressions. But instead of creating some sort of algebraic data type to represent expressions, generating all possible expressions, and evaluating, I sort of interleaved the processes. First, I generate the list 4,3,2,1 with operations optionally applied along the way in all possible ways, then reduce them in all possible ways to a single result, discarding duplicates the whole time. Here’s the code:

import Data.Ratio
import Data.List

type Op = (Rational -> Rational -> Rational)
type Stack = [Rational]

-- make a special kind of division to ignore division by zero. This
-- doesn't give any spurious results since if we can get zero as one
-- of the arguments, we can legitimately create another zero by
-- multiplying.
(//) :: Rational -> Rational -> Rational
a // 0 = 0
a // b = a / b

-- turn a normal binary operator into a function which operates 
-- on the top two elements of a stack.
mkStackOp :: Op -> Stack -> Stack
mkStackOp op (x1:x2:xs) = (x2 `op` x1) : xs
mkStackOp op s = s

negateTop :: Stack -> Stack
negateTop (x:xs) = (negate x) : xs

-- operations that reduce a stack (i.e. binary operations)
stackReducers :: [Stack -> Stack]
stackReducers = map mkStackOp [(+), (-), (*), (//)]

-- operations that transform a stack without reducing it (unary
-- operations).  to allow unary negation, just add negateTop to the
-- list.
stackTransformers :: [Stack -> Stack]
stackTransformers = [id]

allStackOps = stackReducers ++ stackTransformers

-- build up a stack by adding one more element (applying all possible
-- stack transformers), while applying all possible operations to the
-- previous elements.
build :: Rational -> Stack -> [Stack]
build n []  =     [ f [n]             | f  <- stackTransformers ]
build n stk = nub [ f [n] ++ (f' stk) | f  <- stackTransformers,
                                        f' <- allStackOps       ]

-- perform one reduction on a stack in all possible ways.
reduce1 :: Stack -> [Stack]
reduce1 [x] = [[x]]
reduce1 stk = [ f stk | f <- stackReducers ]

-- like >>=, but discarding duplicates.  Ideally we would
--   do this with a Monad instance of Data.Set, but that's
--   currently not possible without doing some contortions
--   to redefine the Monad class (since currently there's no
--   way to define a monad over a subcategory of Haskell types,
--   like we would need to define a Data.Set monad over only
--   Eq types).  See http://www.randomhacks.net/articles/
--   2007/03/15/data-set-monad-haskell-macros.
l >>- f = nub $ concatMap f l

-- completely reduce a stack to a single number in all possible ways.
reduce :: Stack -> [Stack]
reduce [x] = [[x]]
reduce stk = reduce1 stk >>- reduce

-- build up stacks with the given rationals, then reduce.
buildAndReduce :: [Rational] -> Stack -> [Stack]
buildAndReduce [] = reduce
buildAndReduce (r:rs) = s -> (build r s >>- buildAndReduce rs)

-- given a list of starting numbers, return the list of all possible
-- results using arithmetic operators and parentheses on the numbers
-- in the given order.
results :: [Rational] -> [Rational]
results rs = sort $ concat $ buildAndReduce rs []

There’s a little ambiguity in the description of the problem: are we allowed to use – as prefix negation, or just as binary subtraction? The code above treats it only as binary subtraction. In that case we get 52 distinct results ranging from -5 = 4 – 3 * (2+1) to 36 = 4 * 3 * (2+1). 8 are negative, 28 are integers:

Prelude> :l fours
[1 of 1] Compiling Main             ( fours.hs, interpreted )
Ok, modules loaded: Main.
*Main> results [4,3,2,1]
[(-5)%1,(-3)%1,(-2)%1,(-5)%3,(-1)%1,(-2)%3,(-1)%2,(-1)%3,
0%1,1%3,4%9,1%2,4%7,2%3,4%5,1%1,4%3,3%2,8%5,5%3,
2%1,7%3,5%2,8%3,3%1,10%3,7%2,11%3,4%1,13%3,9%2,
5%1,11%2,6%1,13%2,7%1,8%1,9%1,10%1,11%1,12%1,13%1,
14%1,15%1,16%1,20%1,21%1,23%1,24%1,25%1,28%1,36%1]
*Main> length it
52
*Main> length $ filter ((1==) . denominator) $ results [4,3,2,1]
28
*Main> length $ filter (<0) $ results [4,3,2,1]
8

If we allow unary negation, then we get 87 distinct possibilities, 47 of which are integers (and, of course, exactly half of the nonzero possibilities are negative, since they are just the negations of the positive possibilities).