komachi.hs

module Main where

import Data.List
import Data.Maybe
import System.Environment
import Codec.Binary.UTF8.String

data Op = Add | Sub | Mul deriving (Eq, Show)
type OpPos = (Op, Int, Int)
type Answer = [OpPos]

findAnswer :: Int -> [Int] -> Maybe Answer
findAnswer ans nums = listToMaybe $
  mapMaybe (evaluate ans nums)
           [(op, pos1, pos2) | op <- [Add, Sub, Mul], pos1 <- [0..(length nums - 1)],
                               pos2 <- if op == Sub then [0..(length nums - 2)] else [pos1..(length nums - 2)]]

evaluate :: Int -> [Int] -> OpPos -> Maybe Answer
evaluate ans nums (op, pos1, pos2) =
  case nums' of [n] -> if n == ans then Just [(op, pos1, pos2)] else Nothing
                _   -> do res <- findAnswer ans nums'
                          return $ ((op, pos1, pos2):res)
  where (lhs, rhs, rest) = get2Elems pos1 pos2 nums
        nums' = calc op lhs rhs : rest

get2Elems :: Int -> Int -> [a] -> (a, a, [a])
get2Elems pos1 pos2 nums = (n1, n2, hd2 ++ tl2)
  where (hd1, (n1:tl1)) = splitAt pos1 nums
        (hd2, (n2:tl2)) = splitAt pos2 (hd1 ++ tl1)

answerToString :: Int -> [Int] -> Maybe Answer -> String
answerToString _   _    Nothing    = encodeString ""
answerToString ans nums (Just sol) = show ans ++ " = " ++ toExpr (map show nums) sol

toExpr :: [String] -> Answer -> String
toExpr [n] [] = n
toExpr nums ((op, pos1, pos2):sol) = toExpr (join lhs rhs op : rest) sol
  where (lhs, rhs, rest) = get2Elems pos1 pos2 nums
        join n1 n2 op = case op of Add -> "(" ++ n1 ++ "+" ++ n2 ++ ")"
                                   Sub -> "(" ++ n1 ++ "-" ++ n2 ++ ")"
                                   Mul -> n1 ++ "*" ++ n2

calc :: Op -> Int -> Int -> Int
calc Add = (+)
calc Sub = (-)
calc Mul = (*)

main :: IO ()
main = do (ans:nums) <- fmap (map (read :: String -> Int)) getArgs
          putStrLn $ answerToString ans nums $ findAnswer ans nums