public

sztupi / Main.hs

Created 2012-05-23

module Main where
import Prelude
import Control.Applicative
import Data.Maybe
import Control.Monad
import Control.Arrow
import Data.Map as M
import Data.List as L
 
ellie = "Ellie"
emily = "Emily"
jessica = "Jessica"
girls = [ellie, emily, jessica]
 
elephant = "Elephant"
zebra = "Zebra" 
giraffe = "Giraffe"
animals = [elephant, zebra, giraffe]
 
grunge = "Grunge" 
plooper = "Plooper" 
zinger = "Zinger"
lollies = [grunge, plooper, zinger]
 
aunt = "Aunt" 
grandma = "Grandma" 
mum = "Mum"
adults = [aunt, grandma, mum] 
 
type Choices = Map Int [String]
choices = M.fromAscList $ zip [0..] [girls, animals, lollies, adults]
 
type Row = Map Int String
type Table = [Row]
    
main :: IO ()
main = print $ solve
 
(==>) :: (a -> Bool) -> (a -> Bool) -> (a -> Bool)
f ==> g = \x -> (not $ f x) || (g x)
 
have :: Int -> String -> Row -> Bool
have idx value row = (M.lookup idx row) == (Just value)
  
mayHave :: Int -> String -> Row -> Bool
mayHave idx value row = 
    case (M.lookup idx row) of 
      Just v | v /= value -> False
      _ -> True
  
checkRules :: Row -> Bool
checkRules row = 
    all ($row) rules
  where 
    rules = [ (have 0 ellie)    ==> (not . (have 2 zinger))
            , (have 1 giraffe)  ==> (not . (have 2 plooper))
            , (have 3 aunt)     ==> (not . (have 0 jessica)) 
            , (have 0 ellie)    ==> (mayHave 3 grandma)
            , (have 3 mum)      ==> (mayHave 1 elephant)
            , (have 2 grunge)   ==> (mayHave 3 aunt) ]
 
solve :: [Table]
solve = do
    table <- goodTables checkRules choices
    return table
 
isValid :: Choices -> Bool
isValid = L.all (/=[]) . M.elems
 
goodTables :: (Row -> Bool) -> Choices -> [Table]
goodTables check choices = do
    row <- goodRows check choices
    case remains choices row of
      remainingChoices | (isValid remainingChoices) -> do
        rows <- goodTables check remainingChoices
        return $ row:rows
      _ -> return [row]
 
goodRows :: (Row -> Bool) -> Choices -> [Row]
goodRows check choices = 
    goodRows' (M.empty) 0 
  where
    goodRows' :: Row -> Int -> [Row]
    goodRows' start idx | idx == M.size choices = [start]
    goodRows' start idx = do
        pick <- choices ! idx
        let row = M.insert idx pick start
        guard $ check row
        goodRows' row (idx+1) 
 
remains :: Choices -> Row -> Choices 
remains original row =
    M.foldWithKey copyWithoutRowValue M.empty original
  where
    copyWithoutRowValue :: Int -> [String] -> Choices -> Choices
    copyWithoutRowValue key value acc = M.insert key (L.delete (row ! key) value) acc