send + more = money

send-more-money.hs

{-# LANGUAGE TupleSections #-}

module Main where

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Trans.State
import           Data.List                 ((\\), transpose)
import           Data.Map                  (Map)
import qualified Data.Map                  as M
import           Data.Set                  (Set)
import qualified Data.Set                  as S
import           System.Environment

-- Read arguments from command line and solve.
main :: IO ()
main = getArgs >>= print . solveCryptarithm

-- -----------------------------------------------------------------------------
-- Cryptarithm API

type Digit  = Int
type Digits = [Int]
type Sum    = Int

type Cryptarithm = StateT (Digits, Sum) []

runCryptarithm :: Cryptarithm a -> [a]
runCryptarithm = flip evalStateT ([0..9], 0)

-- Select a new summand digit.
selectSummand :: Cryptarithm Digit
selectSummand = StateT go
  where
    go :: (Digits, Sum) -> [(Digit, (Digits, Sum))]
    go ([], _) = []
    go (x:xs, sum) = (x, (xs, sum + x)) : [ (y, (x:ys, sum')) | (y, (ys, sum')) <- go (xs, sum) ]

-- Select a new sum digit.
selectSum :: Cryptarithm Digit
selectSum = do
    (xs, carry, x) <- getSplit
    case find x xs of
        (Nothing, _) -> mzero
        (Just _, xs') -> do
            put (xs', carry)
            return x
  where
    find :: Eq a => a -> [a] -> (Maybe a, [a])
    find = go []
      where
        go acc _ [] = (Nothing, acc)
        go acc x (y:ys)
          | x == y = (Just y, acc ++ ys)
          | otherwise = go (y:acc) x ys

-- Use an already-selected sum digit.
useSum :: Digit -> Cryptarithm ()
useSum x = do
    (xs, carry, x') <- getSplit
    guard (x == x')
    put (xs, carry)

-- Use an already-selected summand digit.
useSummand :: Digit -> Cryptarithm ()
useSummand x = modify (\(xs, s) -> (xs, x + s))

-- Modify a Cryptarithm computation to only select a non-zero digit.
nonZero :: Cryptarithm Digit -> Cryptarithm Digit
nonZero action = do
    x <- action
    guard (x /= 0)
    return x

-- ------------
-- Helper funcs

getSplit :: StateT (Digits, Sum) [] (Digits, Sum, Digit)
getSplit = do
    (xs, sum) <- get
    let (carry, x) = divMod sum 10
    return (xs, carry, x)

val :: Digits -> Int
val = foldl (\acc x -> acc*10 + x) 0

-- -----------------------------------------------------------------------------
-- Examples of how to solve "by hand"

--   S E N D
-- + M O R E
-- ---------
-- M O N E Y
sendMoreMoney :: [(Int, Int, Int)]
sendMoreMoney = runCryptarithm $ do
    d <- selectSummand
    e <- selectSummand
    y <- selectSum
    n <- selectSummand
    r <- selectSummand
    _ <- useSum e
    _ <- useSummand e
    o <- selectSummand
    _ <- useSum n
    s <- nonZero selectSummand
    m <- nonZero selectSummand
    _ <- useSum o
    _ <- useSum m

    let send  = val [s,e,n,d]
        more  = val [m,o,r,e]
        money = val [m,o,n,e,y]

    return (send, more, money)

--       B I L L
-- W I L L I A M
--   M O N I C A
-- -------------
-- C L I N T O N
billWilliamMonicaClinton = runCryptarithm $ do
    l <- selectSummand
    m <- nonZero selectSummand
    a <- selectSummand
    n <- selectSum
    _ <- useSummand l
    _ <- useSummand a
    c <- selectSummand
    o <- selectSum
    i <- selectSummand
    _ <- useSummand i
    _ <- useSummand i
    t <- selectSum
    b <- nonZero selectSummand
    _ <- useSummand l
    _ <- useSummand n
    _ <- useSum n
    _ <- useSummand l
    _ <- useSummand o
    _ <- useSum i
    _ <- useSummand i
    _ <- useSummand m
    _ <- useSum l
    w <- nonZero selectSummand
    _ <- useSum c

    let bill    = val [b,i,l,l]
        william = val [w,i,l,l,i,a,m]
        monica  = val [m,o,n,i,c,a]
        clinton = val [c,l,i,n,t,o,n]

    return (bill, william, monica, clinton)

--     D O S
--     D O S
--   T R E S
-- ---------
-- S I E T E
--
dosDosTresSiete :: [(Int, Int, Int)]
dosDosTresSiete = runCryptarithm $ do
    s <- nonZero selectSummand
    _ <- useSummand s
    _ <- useSummand s
    e <- selectSum
    o <- selectSummand
    _ <- useSummand o
    _ <- useSummand e
    t <- nonZero selectSum
    d <- nonZero selectSummand
    _ <- useSummand d
    r <- selectSummand
    _ <- useSum e
    _ <- useSummand t
    i <- selectSum
    _ <- useSum s

    let dos   = val [d,o,s]
        tres  = val [t,r,e,s]
        siete = val [s,i,e,t,e]

    return (dos, tres, siete)

-- -----------------------------------------------------------------------------
-- String-based solver (["foo","bar","baz"] means "solve for: foo + bar = baz")

solveCryptarithm :: [String] -> [[Int]]
solveCryptarithm xs = go (S.fromList (map head xs)) xs
  where
    go :: Set Char -> [String] -> [[Int]]
    go firsts = map (makeInts xs)
              . runCryptarithm
              . foldM solveOneColumn M.empty
              . reverse
              . transpose
              . padLeft
      where
        -- Pad a list of strings with spaces on the left so that
        -- each string is the length of the last (assumed to be
        -- the longest or tied for the longest).
        padLeft :: [String] -> [String]
        padLeft [] = []
        padLeft xs = map (padTo (length (last xs))) xs
          where
            padTo :: Int -> String -> String
            padTo n s = let len = length s
                        in if len == n
                               then s
                               else replicate (n - len) ' ' ++ s

        solveOneColumn :: Map Char Digit -> String -> Cryptarithm (Map Char Digit)
        solveOneColumn m [c] =
            case M.lookup c m of
                Just digit -> useSum digit >> return m
                Nothing    -> (\digit -> M.insert c digit m) <$> if S.member c firsts
                                                                     then nonZero selectSum
                                                                     else selectSum
        solveOneColumn m (' ':cs) = solveOneColumn m cs
        solveOneColumn m (c:cs) =
            case M.lookup c m of
                Just digit -> useSummand digit >> solveOneColumn m cs
                Nothing    -> (if S.member c firsts
                                   then nonZero selectSummand
                                   else selectSummand) >>=
                              \digit -> solveOneColumn (M.insert c digit m) cs

        makeInts :: [String] -> Map Char Digit -> [Int]
        makeInts xs m = map (makeInt m) xs
          where
            makeInt :: Map Char Digit -> String -> Int
            makeInt m = val . map (m M.!)