twentyeighteen.hs

module Main where

import Control.Monad
import Data.List (foldl')

data Operator
  = Add
  | Sub
  | Mult
  | Div
  deriving (Show, Eq)

data Token
  = Number !Rational
  | Op !Operator
  deriving (Show, Eq)

type RpnTokens = [Token]

type Stack a = [a]

eval :: RpnTokens -> Maybe Rational
eval = fmap head . foldM f []
  where
    f stack (Number n) = Just $ (n : stack)
    f (b:a:stack) (Op op) =
      case op of
        Add -> Just $ (a + b) : stack
        Sub -> Just $ (a - b) : stack
        Mult -> Just $ (a * b) : stack
        Div ->
          if b == 0
            then Nothing
            else Just $ (a / b) : stack

display :: RpnTokens -> String
display = head . foldl' go [""]
  where
    go stack (Number n) = (show . (round :: Rational -> Int) $ n) : stack
    go (b:a:stack) (Op op) =
      case op of
        Add -> ("(" ++ a ++ "+" ++ b ++ ")") : stack
        Sub -> ("(" ++ a ++ "-" ++ b ++ ")") : stack
        Mult -> ("(" ++ a ++ "*" ++ b ++ ")") : stack
        Div -> ("(" ++ a ++ "/" ++ b ++ ")") : stack

generate :: [Int] -> [[Token]]
generate numbers = f numbers 0 0
  where
    totalLength = 2 * (length numbers) - 1
    operators = [Op Add, Op Sub, Op Mult, Op Div]
    f ns numNumbers numOps
      | numNumbers + numOps == totalLength = [[]]
      | numNumbers > numOps + 1 =
        if null ns
          then opNext
          else numberNext ++ opNext
      | otherwise = numberNext
      where
        n' = Number . fromIntegral . head $ ns
        numberNext = (:) <$> [n'] <*> f (tail ns) (numNumbers + 1) numOps
        opNext = (:) <$> operators <*> f ns numNumbers (numOps + 1)

main :: IO ()
main = do
  let solutions =
        filter ((== Just 2018) . eval) . generate . reverse $ [1 .. 10]
  -- Show any solution
  putStrLn . display . head $ solutions
  -- Enumerate all solutions: 
  -- forM_ (zip [1 ..] solutions) $ \(n, tokens) ->
  --   putStrLn $ (show n) ++ ": " ++ display tokens