Haskell Quiz/Verbal Arithmetic/Solution Dolio

This solution doesn't do anything particularly special. It parses the equations using a simple Parsec parser, and then searches for solutions using backtracking search via the list monad. However, to make things a bit simpler, I wrapped up the searching into what I called the CSP monad (for constraint satisfaction problem), that automatically handles the remembering of already assigned variables, and branching when attempting to use a previously unassigned variable. It's not a clever solution, algorithmically, but it gets the job done for such a small problem as this.

The problem specific code:

module Main(main) where

import Control.Monad

import Data.Char
import Data.List
import Data.Maybe

import Data.Map (Map)
import qualified Data.Map as M

import Text.ParserCombinators.Parsec

import CSP

data Equation = Equation { summands :: [String], result :: String } deriving Show

-- A quick parser for the equations in use. I didn't bother trying to cover all
-- eventualities, but it will parse 'word1 + word2 + word3 ... = result'
add = char '+' >> spaces >> return (++)
equ = char '=' >> spaces
vnum = many1 lower >>= \a -> spaces >> return a

equation = do l <- chainl1 (liftM return vnum) add
              equ
              r <- vnum
              return (Equation l r)

-- Constraints on individual digits. Digits in a column, plus the carry-in, mod 10
-- must equal the result digit. Returns the carry-out
digitcsp :: [Char] -> Int -> Char -> CSP Int Int
digitcsp vs carry result = do l <- mapM byName vs
                              let (d, m) = divMod (carry + sum l) 10
                              byName result >>= guard . (m ==)
                              return d

-- Constraint to make sure the left-most digit on any term is not zero. We don't want
-- 'two = 012', for instance.
sigfig :: Equation -> CSP Int ()
sigfig (Equation ss r) = chk r >> foldM_ (const chk) () ss
 where
 chk (h:_) = byName h >>= guard . (0 /=)

-- The overall constraint for solving equations
eqncsp :: Equation -> CSP Int ()
eqncsp e@(Equation s r@(hr:_)) = go (columns s) (reverse r) 0 >> sigfig e
 where
 go (vs:vss) (r:rs) carry = digitcsp vs carry r >>= go vss rs
 go []       [r]    carry = byName r >>= guard . (carry ==)
 go []       []     0     = return ()
 go _        _      _     = mzero

-- turns a list of terms into a list of letters in each column, to be summed.
columns :: [String] -> [[Char]]
columns = transpose . map reverse

-- Given an equation, and mappings for all the letters, constructs the numerals
showEqn :: Equation -> Map Char Int -> String
showEqn (Equation ss r) m = intercalate " + " (map showWord ss) ++ " = " ++ showWord r
 where
 showWord = map (\d -> intToDigit (fromJust (M.lookup d m)))

main = mapM_ (display . fmap solve . parse equation "foo") . lines =<< getContents
 where
 display (Left _) = putStrLn "Bad parse."
 display (Right ss) = mapM_ putStrLn ss

solve :: Equation -> [String]
solve e = map (showEqn e) ms
 where ms = execCSP (eqncsp e) [0..9]

And the CSP monad. There is an auxiliary 'SupplyT' monad which encapsulates selecting non-deterministically from a stored supply.

{-# OPTIONS_GHC -fglasgow-exts #-}

module CSP where

import Data.Map (Map)
import qualified Data.Map as M

import Control.Arrow
import Control.Monad
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Trans

newtype CSP b a = CSP { unCSP :: StateT (Map Char b) (SupplyT b []) a }
                    deriving (Monad, MonadPlus, MonadState (Map Char b),
                                MonadSupply b)

newtype SupplyT s m a = SupplyT { unSupplyT :: StateT [s] m a }
                        deriving (Monad, MonadPlus, MonadState [s])

class MonadPlus m => MonadSupply s m where
    getSupply :: m s

instance MonadPlus m => MonadSupply s (SupplyT s m) where
    getSupply = do l <- get
                   (v, l') <- choose (pick l)
                   put l'
                   return v
     where
     pick (h:t) = (h, t) : map (second (h:)) (pick t)
     pick []    = mzero

instance MonadSupply s m => MonadSupply s (StateT r m) where
    getSupply = lift getSupply

runSupplyT :: Monad m => SupplyT s m a -> [s] -> m a
runSupplyT st s = flip evalStateT s . unSupplyT $ st

runCSP :: CSP b a -> [b] -> [(a, Map Char b)]
runCSP c l = flip runSupplyT l . flip runStateT M.empty . unCSP $ c

evalCSP :: CSP b a -> [b] -> [a]
evalCSP c l = flip runSupplyT l . flip evalStateT M.empty . unCSP $ c

execCSP :: CSP b a -> [b] -> [Map Char b]
execCSP c l = flip runSupplyT l . flip execStateT M.empty . unCSP $ c

byName :: Char -> CSP b b
byName k = gets (M.lookup k) >>= maybe c return
 where c = getSupply >>= assign k

assign :: Char -> a -> CSP a a
assign k a = modify (M.insert k a) >> return a

choose :: MonadPlus m => [a] -> m a
choose = msum . map return