{- |
Description: An arithmetic puzzle
Copyright: Copyright (C) 2023 Yoo Chung
License: GPL-3.0-or-later
Maintainer: dev@chungyc.org

Some solutions to "Problems.P93" of Ninety-Nine Haskell "Problems".
-}
module Solutions.P93 (arithmeticPuzzle) where

import           Data.List            (inits, nub, tails)
import           Data.Map.Lazy        (Map, (!))
import qualified Data.Map.Lazy        as Map
import           Data.Maybe           (mapMaybe)
import           Data.Ratio
import qualified Data.Set             as Set
import           Problems.BinaryTrees

{- |
Given a list of integer numbers, find a correct way of inserting
the arithmetic signs such that the result is a correct equation.
-}
arithmeticPuzzle :: [Integer] -> [String]
arithmeticPuzzle :: [Integer] -> [String]
arithmeticPuzzle []  = []
arithmeticPuzzle [Integer
_] = []
arithmeticPuzzle [Integer]
xs  = [String] -> [String]
forall a. Eq a => [a] -> [a]
nub ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ (Tree (Either Op Integer) -> String)
-> [Tree (Either Op Integer)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Tree (Either Op Integer) -> String
formatEquation [Tree (Either Op Integer)]
es
  where ts :: [Tree (Maybe Integer)]
ts = [Integer] -> [Tree (Maybe Integer)]
toTrees [Integer]
xs
        es :: [Tree (Either Op Integer)]
es = (Tree (Maybe Integer) -> [Tree (Either Op Integer)])
-> [Tree (Maybe Integer)] -> [Tree (Either Op Integer)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree (Maybe Integer) -> [Tree (Either Op Integer)]
findEquations [Tree (Maybe Integer)]
ts

-- | Form all binary trees that can be formed from the list of integers,
-- where all internal nodes are 'Nothing', which is the placeholder for arithmetic operations,
-- and all leaves are numbers.  All internal nodes have two non-empty subtrees.
toTrees :: [Integer] -> [Tree (Maybe Integer)]
toTrees :: [Integer] -> [Tree (Maybe Integer)]
toTrees []  = [Tree (Maybe Integer)
forall a. Tree a
Empty]
toTrees [Integer
x] = [Maybe Integer
-> Tree (Maybe Integer)
-> Tree (Maybe Integer)
-> Tree (Maybe Integer)
forall a. a -> Tree a -> Tree a -> Tree a
Branch (Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
x) Tree (Maybe Integer)
forall a. Tree a
Empty Tree (Maybe Integer)
forall a. Tree a
Empty]
toTrees [Integer]
xs  = (([Integer], [Integer]) -> [Tree (Maybe Integer)])
-> [([Integer], [Integer])] -> [Tree (Maybe Integer)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\([Integer]
ls,[Integer]
rs) -> [Maybe Integer
-> Tree (Maybe Integer)
-> Tree (Maybe Integer)
-> Tree (Maybe Integer)
forall a. a -> Tree a -> Tree a -> Tree a
Branch Maybe Integer
forall a. Maybe a
Nothing Tree (Maybe Integer)
l Tree (Maybe Integer)
r | Tree (Maybe Integer)
l <- [Integer] -> [Tree (Maybe Integer)]
toTrees [Integer]
ls, Tree (Maybe Integer)
r <- [Integer] -> [Tree (Maybe Integer)]
toTrees [Integer]
rs]) [([Integer], [Integer])]
splits
  where splits :: [([Integer], [Integer])]
splits = (([Integer], [Integer]) -> Bool)
-> [([Integer], [Integer])] -> [([Integer], [Integer])]
forall a. (a -> Bool) -> [a] -> [a]
filter (\([Integer]
ls,[Integer]
rs) -> Bool -> Bool
not ([Integer] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Integer]
ls) Bool -> Bool -> Bool
&& Bool -> Bool
not ([Integer] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Integer]
rs)) ([([Integer], [Integer])] -> [([Integer], [Integer])])
-> [([Integer], [Integer])] -> [([Integer], [Integer])]
forall a b. (a -> b) -> a -> b
$ [[Integer]] -> [[Integer]] -> [([Integer], [Integer])]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Integer] -> [[Integer]]
forall a. [a] -> [[a]]
inits [Integer]
xs) ([Integer] -> [[Integer]]
forall a. [a] -> [[a]]
tails [Integer]
xs)

-- | Arithmetic operations which are internal node values for an expression tree.
data Op = Add | Subtract | Multiply | Divide | Equals

-- | Find all equations that can be formed from the binary tree structure.
-- Both sides must be equal to each other.
findEquations :: Tree (Maybe Integer) -> [Tree (Either Op Integer)]
findEquations :: Tree (Maybe Integer) -> [Tree (Either Op Integer)]
findEquations (Branch Maybe Integer
Nothing Tree (Maybe Integer)
l Tree (Maybe Integer)
r) = (Ratio Integer -> [Tree (Either Op Integer)])
-> [Ratio Integer] -> [Tree (Either Op Integer)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Ratio Integer -> [Tree (Either Op Integer)]
toEquations [Ratio Integer]
common
  where ls :: Map (Ratio Integer) [Tree (Either Op Integer)]
ls = Tree (Maybe Integer)
-> Map (Ratio Integer) [Tree (Either Op Integer)]
findExpressions Tree (Maybe Integer)
l
        rs :: Map (Ratio Integer) [Tree (Either Op Integer)]
rs = Tree (Maybe Integer)
-> Map (Ratio Integer) [Tree (Either Op Integer)]
findExpressions Tree (Maybe Integer)
r
        common :: [Ratio Integer]
common = Set (Ratio Integer) -> [Ratio Integer]
forall a. Set a -> [a]
Set.toList (Set (Ratio Integer) -> [Ratio Integer])
-> Set (Ratio Integer) -> [Ratio Integer]
forall a b. (a -> b) -> a -> b
$ Set (Ratio Integer) -> Set (Ratio Integer) -> Set (Ratio Integer)
forall a. Ord a => Set a -> Set a -> Set a
Set.intersection (Map (Ratio Integer) [Tree (Either Op Integer)]
-> Set (Ratio Integer)
forall k a. Map k a -> Set k
Map.keysSet Map (Ratio Integer) [Tree (Either Op Integer)]
ls) (Map (Ratio Integer) [Tree (Either Op Integer)]
-> Set (Ratio Integer)
forall k a. Map k a -> Set k
Map.keysSet Map (Ratio Integer) [Tree (Either Op Integer)]
rs)
        toEquations :: Ratio Integer -> [Tree (Either Op Integer)]
toEquations Ratio Integer
n = [Either Op Integer
-> Tree (Either Op Integer)
-> Tree (Either Op Integer)
-> Tree (Either Op Integer)
forall a. a -> Tree a -> Tree a -> Tree a
Branch (Op -> Either Op Integer
forall a b. a -> Either a b
Left Op
Equals) Tree (Either Op Integer)
lt Tree (Either Op Integer)
rt | Tree (Either Op Integer)
lt <- Map (Ratio Integer) [Tree (Either Op Integer)]
ls Map (Ratio Integer) [Tree (Either Op Integer)]
-> Ratio Integer -> [Tree (Either Op Integer)]
forall k a. Ord k => Map k a -> k -> a
! Ratio Integer
n, Tree (Either Op Integer)
rt <- Map (Ratio Integer) [Tree (Either Op Integer)]
rs Map (Ratio Integer) [Tree (Either Op Integer)]
-> Ratio Integer -> [Tree (Either Op Integer)]
forall k a. Ord k => Map k a -> k -> a
! Ratio Integer
n]
findEquations Tree (Maybe Integer)
_ = [Tree (Either Op Integer)]
forall a. HasCallStack => a
undefined  -- need two sides for an equation

-- | Find all expression trees that can be formed from the given binary tree structure.
-- Return a map from their values to the trees.
findExpressions :: Tree (Maybe Integer) -> Map (Ratio Integer) [Tree (Either Op Integer)]
findExpressions :: Tree (Maybe Integer)
-> Map (Ratio Integer) [Tree (Either Op Integer)]
findExpressions Tree (Maybe Integer)
t = ([Tree (Either Op Integer)]
 -> [Tree (Either Op Integer)] -> [Tree (Either Op Integer)])
-> [(Ratio Integer, [Tree (Either Op Integer)])]
-> Map (Ratio Integer) [Tree (Either Op Integer)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [Tree (Either Op Integer)]
-> [Tree (Either Op Integer)] -> [Tree (Either Op Integer)]
forall a. [a] -> [a] -> [a]
(++) ([(Ratio Integer, [Tree (Either Op Integer)])]
 -> Map (Ratio Integer) [Tree (Either Op Integer)])
-> [(Ratio Integer, [Tree (Either Op Integer)])]
-> Map (Ratio Integer) [Tree (Either Op Integer)]
forall a b. (a -> b) -> a -> b
$ (Tree (Either Op Integer)
 -> Maybe (Ratio Integer, [Tree (Either Op Integer)]))
-> [Tree (Either Op Integer)]
-> [(Ratio Integer, [Tree (Either Op Integer)])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\Tree (Either Op Integer)
t' -> (Maybe (Ratio Integer), Tree (Either Op Integer))
-> Maybe (Ratio Integer, [Tree (Either Op Integer)])
forall {a} {a}. (Maybe a, a) -> Maybe (a, [a])
assoc (Tree (Either Op Integer) -> Maybe (Ratio Integer)
evalTree Tree (Either Op Integer)
t', Tree (Either Op Integer)
t')) ([Tree (Either Op Integer)]
 -> [(Ratio Integer, [Tree (Either Op Integer)])])
-> [Tree (Either Op Integer)]
-> [(Ratio Integer, [Tree (Either Op Integer)])]
forall a b. (a -> b) -> a -> b
$ Tree (Maybe Integer) -> [Tree (Either Op Integer)]
permuteTree Tree (Maybe Integer)
t
  where assoc :: (Maybe a, a) -> Maybe (a, [a])
assoc (Maybe a
Nothing, a
_) = Maybe (a, [a])
forall a. Maybe a
Nothing
        assoc (Just a
x, a
t') = (a, [a]) -> Maybe (a, [a])
forall a. a -> Maybe a
Just (a
x, [a
t'])

-- | Compute the value of an expression tree.
-- Returns 'Nothing' if it is invalid, e.g., there would be division by zero.
evalTree :: Tree (Either Op Integer) -> Maybe (Ratio Integer)
evalTree :: Tree (Either Op Integer) -> Maybe (Ratio Integer)
evalTree (Branch (Right Integer
n) Tree (Either Op Integer)
Empty Tree (Either Op Integer)
Empty) = Ratio Integer -> Maybe (Ratio Integer)
forall a. a -> Maybe a
Just (Ratio Integer -> Maybe (Ratio Integer))
-> Ratio Integer -> Maybe (Ratio Integer)
forall a b. (a -> b) -> a -> b
$ Integer -> Ratio Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n

evalTree (Branch (Left Op
Divide) Tree (Either Op Integer)
l Tree (Either Op Integer)
r) = do
  Ratio Integer
l' <- Tree (Either Op Integer) -> Maybe (Ratio Integer)
evalTree Tree (Either Op Integer)
l
  Ratio Integer
r' <- Tree (Either Op Integer) -> Maybe (Ratio Integer)
evalTree Tree (Either Op Integer)
r
  case Ratio Integer
r' of Ratio Integer
0 -> Maybe (Ratio Integer)
forall a. Maybe a
Nothing
             Ratio Integer
_ -> Ratio Integer -> Maybe (Ratio Integer)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Ratio Integer -> Maybe (Ratio Integer))
-> Ratio Integer -> Maybe (Ratio Integer)
forall a b. (a -> b) -> a -> b
$ Ratio Integer
l' Ratio Integer -> Ratio Integer -> Ratio Integer
forall a. Fractional a => a -> a -> a
/ Ratio Integer
r'

evalTree (Branch (Left Op
op) Tree (Either Op Integer)
l Tree (Either Op Integer)
r) = do
  Ratio Integer
l' <- Tree (Either Op Integer) -> Maybe (Ratio Integer)
evalTree Tree (Either Op Integer)
l
  Ratio Integer
r' <- Tree (Either Op Integer) -> Maybe (Ratio Integer)
evalTree Tree (Either Op Integer)
r
  Ratio Integer -> Maybe (Ratio Integer)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Ratio Integer -> Maybe (Ratio Integer))
-> Ratio Integer -> Maybe (Ratio Integer)
forall a b. (a -> b) -> a -> b
$ Op -> Ratio Integer -> Ratio Integer -> Ratio Integer
evaluate Op
op Ratio Integer
l' Ratio Integer
r'

evalTree Tree (Either Op Integer)
_                              = Maybe (Ratio Integer)
forall a. HasCallStack => a
undefined

evaluate :: Op -> Ratio Integer -> Ratio Integer -> Ratio Integer
evaluate :: Op -> Ratio Integer -> Ratio Integer -> Ratio Integer
evaluate Op
Add Ratio Integer
x Ratio Integer
y      = Ratio Integer
x Ratio Integer -> Ratio Integer -> Ratio Integer
forall a. Num a => a -> a -> a
+ Ratio Integer
y
evaluate Op
Subtract Ratio Integer
x Ratio Integer
y = Ratio Integer
x Ratio Integer -> Ratio Integer -> Ratio Integer
forall a. Num a => a -> a -> a
- Ratio Integer
y
evaluate Op
Multiply Ratio Integer
x Ratio Integer
y = Ratio Integer
x Ratio Integer -> Ratio Integer -> Ratio Integer
forall a. Num a => a -> a -> a
* Ratio Integer
y
evaluate Op
Divide Ratio Integer
x Ratio Integer
y   = Ratio Integer
x Ratio Integer -> Ratio Integer -> Ratio Integer
forall a. Fractional a => a -> a -> a
/ Ratio Integer
y
evaluate Op
_ Ratio Integer
_ Ratio Integer
_        = Ratio Integer
forall a. HasCallStack => a
undefined

permuteTree :: Tree (Maybe Integer) -> [Tree (Either Op Integer)]
permuteTree :: Tree (Maybe Integer) -> [Tree (Either Op Integer)]
permuteTree (Branch Maybe Integer
Nothing Tree (Maybe Integer)
l Tree (Maybe Integer)
r)          = [Either Op Integer
-> Tree (Either Op Integer)
-> Tree (Either Op Integer)
-> Tree (Either Op Integer)
forall a. a -> Tree a -> Tree a -> Tree a
Branch (Op -> Either Op Integer
forall a b. a -> Either a b
Left Op
op) Tree (Either Op Integer)
l' Tree (Either Op Integer)
r' | Op
op <- [Op]
ops, Tree (Either Op Integer)
l' <- [Tree (Either Op Integer)]
ls, Tree (Either Op Integer)
r' <- [Tree (Either Op Integer)]
rs]
  where ops :: [Op]
ops = [Op
Add, Op
Subtract, Op
Multiply, Op
Divide]
        ls :: [Tree (Either Op Integer)]
ls = Tree (Maybe Integer) -> [Tree (Either Op Integer)]
permuteTree Tree (Maybe Integer)
l
        rs :: [Tree (Either Op Integer)]
rs = Tree (Maybe Integer) -> [Tree (Either Op Integer)]
permuteTree Tree (Maybe Integer)
r
permuteTree (Branch (Just Integer
n) Tree (Maybe Integer)
Empty Tree (Maybe Integer)
Empty) = [Either Op Integer
-> Tree (Either Op Integer)
-> Tree (Either Op Integer)
-> Tree (Either Op Integer)
forall a. a -> Tree a -> Tree a -> Tree a
Branch (Integer -> Either Op Integer
forall a b. b -> Either a b
Right Integer
n) Tree (Either Op Integer)
forall a. Tree a
Empty Tree (Either Op Integer)
forall a. Tree a
Empty]
permuteTree Tree (Maybe Integer)
_                             = [Tree (Either Op Integer)]
forall a. HasCallStack => a
undefined  -- should not be possible

-- | Returns the given expression tree in string form.
-- It only inserts parentheses which are necessary.
formatEquation :: Tree (Either Op Integer) -> String
formatEquation :: Tree (Either Op Integer) -> String
formatEquation (Branch (Right Integer
n) Tree (Either Op Integer)
Empty Tree (Either Op Integer)
Empty) = Integer -> String
forall a. Show a => a -> String
show Integer
n
formatEquation (Branch (Left Op
Equals) Tree (Either Op Integer)
l Tree (Either Op Integer)
r) = Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" = " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
r
formatEquation (Branch (Left Op
Add) Tree (Either Op Integer)
l Tree (Either Op Integer)
r) = Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"+" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
r

formatEquation (Branch (Left Op
Subtract) Tree (Either Op Integer)
l Tree (Either Op Integer)
r)
  | Tree (Either Op Integer) -> Bool
isAdd Tree (Either Op Integer)
r Bool -> Bool -> Bool
|| Tree (Either Op Integer) -> Bool
isSubtract Tree (Either Op Integer)
r = Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"-(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
r String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
  | Bool
otherwise               = Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"-" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
r

formatEquation (Branch (Left Op
Multiply) Tree (Either Op Integer)
l Tree (Either Op Integer)
r) = Tree (Either Op Integer) -> String
format Tree (Either Op Integer)
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"*" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree (Either Op Integer) -> String
format Tree (Either Op Integer)
r
  where format :: Tree (Either Op Integer) -> String
format Tree (Either Op Integer)
t | Tree (Either Op Integer) -> Bool
isAdd Tree (Either Op Integer)
t Bool -> Bool -> Bool
|| Tree (Either Op Integer) -> Bool
isSubtract Tree (Either Op Integer)
t = String
"(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
t String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
                 | Bool
otherwise               = Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
t

formatEquation (Branch (Left Op
Divide) Tree (Either Op Integer)
l Tree (Either Op Integer)
r)
  | Tree (Either Op Integer) -> Bool
isMultiply Tree (Either Op Integer)
r Bool -> Bool -> Bool
|| Tree (Either Op Integer) -> Bool
isDivide Tree (Either Op Integer)
r = Tree (Either Op Integer) -> String
format Tree (Either Op Integer)
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"/(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
r String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
  | Bool
otherwise                  = Tree (Either Op Integer) -> String
format Tree (Either Op Integer)
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"/" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree (Either Op Integer) -> String
format Tree (Either Op Integer)
r
  where format :: Tree (Either Op Integer) -> String
format Tree (Either Op Integer)
t | Tree (Either Op Integer) -> Bool
isAdd Tree (Either Op Integer)
t Bool -> Bool -> Bool
|| Tree (Either Op Integer) -> Bool
isSubtract Tree (Either Op Integer)
t = String
"(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
t String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
                 | Bool
otherwise               = Tree (Either Op Integer) -> String
formatEquation Tree (Either Op Integer)
t

formatEquation Tree (Either Op Integer)
_ = String
forall a. HasCallStack => a
undefined  -- should not be possible

isAdd :: Tree (Either Op Integer) -> Bool
isAdd :: Tree (Either Op Integer) -> Bool
isAdd (Branch (Left Op
Add) Tree (Either Op Integer)
_ Tree (Either Op Integer)
_) = Bool
True
isAdd Tree (Either Op Integer)
_                       = Bool
False

isSubtract :: Tree (Either Op Integer) -> Bool
isSubtract :: Tree (Either Op Integer) -> Bool
isSubtract (Branch (Left Op
Subtract) Tree (Either Op Integer)
_ Tree (Either Op Integer)
_) = Bool
True
isSubtract Tree (Either Op Integer)
_                            = Bool
False

isMultiply :: Tree (Either Op Integer) -> Bool
isMultiply :: Tree (Either Op Integer) -> Bool
isMultiply (Branch (Left Op
Multiply) Tree (Either Op Integer)
_ Tree (Either Op Integer)
_) = Bool
True
isMultiply Tree (Either Op Integer)
_                            = Bool
False

isDivide ::  Tree (Either Op Integer) -> Bool
isDivide :: Tree (Either Op Integer) -> Bool
isDivide (Branch (Left Op
Divide) Tree (Either Op Integer)
_ Tree (Either Op Integer)
_) = Bool
True
isDivide Tree (Either Op Integer)
_                          = Bool
False