/r/dailyprogrammer #201 hard solution

Solve.hs

import Control.Applicative
import Control.Monad
import Data.Char
import Data.List
import Data.Maybe
import Data.Ord
import Debug.Trace
import Text.Parsec.Char
import Text.Parsec.Combinator hiding (many, (<|>))
import Text.Parsec.Prim hiding (many, (<|>))
import Text.Parsec.String

--
-- data types
--
type ID = String
type Probability = Maybe Int

data ConstraintTensor = CTCons ID ConstraintTensor ConstraintTensor |
                        CTEnd Probability
    deriving (Show)

data ConstraintIndex = Explicit ID |
                       Inverse ID |
                       Both ID
    deriving (Show)

type ConstraintQuery = [ConstraintIndex]

newtype SortedConstraintQuery = SortedConstraintQuery ConstraintQuery deriving (Show)

type QueryResult = (SortedConstraintQuery, Probability)

data Constraint = Sum [SortedConstraintQuery] Probability deriving (Show)

data Problem = Problem {
    problemIDs :: [ID],
    problemSize :: Int, -- equivalent to length . problemIDs
    problemConstraints :: [Constraint],
    problemSpec :: [ConstraintIndex]
} deriving (Show)

--
-- solution logic
--

countMatch pred = length . filter pred

unknownProbability :: QueryResult -> Bool
unknownProbability = isNothing . snd

implicitConstraints :: [ID] -> [Constraint]
implicitConstraints = map implicitConstraint
    where implicitConstraint x = Sum [SortedConstraintQuery [Both x]] (Just 100)

inverseQueryList :: SortedConstraintQuery -> [SortedConstraintQuery]
inverseQueryList (SortedConstraintQuery [])     = []
inverseQueryList (SortedConstraintQuery (x:xs)) =
    SortedConstraintQuery (ix:xs) :
        map (app x) invs ++ map (app ix) invs
    where invs = map (\(SortedConstraintQuery x) -> x) $ inverseQueryList $ SortedConstraintQuery xs
          ix = case x of
                    Explicit i -> Inverse i
                    Inverse i -> Explicit i
                    Both i -> error "Inverse of Both is nonsense"
          app v ys = SortedConstraintQuery (v:ys)

inverseConstraint :: Constraint -> Constraint
inverseConstraint (Sum xs (Just p)) = Sum (concatMap inverseQueryList xs) (Just (100 - p))

resolveConstraint :: Constraint -> ConstraintTensor -> ConstraintTensor
resolveConstraint s@(Sum scq (Just p)) c
    | nf == 0   = c
    | nf == 1   = set c (fst . head $ fixable) $ Just (p - t)
    | otherwise = c --error $ "constraint " ++ show s ++ " not specific enough for " ++ show c ++ " results " ++ show qr
    where t = sum (map (fromJust . snd) . filter (isJust . snd) $ qr)
          fixable = filter unknownProbability qr
          nf = length fixable
          qr = concatMap (c !) scq

sortConstraints :: ConstraintTensor -> [Constraint] -> [Constraint]
sortConstraints c = sortBy (comparing $ countMatch unknownProbability . concatMap (c!) . sumQueries)

resolveConstraints :: [Constraint] -> ConstraintTensor -> ConstraintTensor
resolveConstraints []     c = c
resolveConstraints ys c = seq c' $ resolveConstraints xs c'
    where c'     = resolveConstraint x c
          (x:xs) = sortConstraints c ys

main :: IO()
main = do
    s <- parseFromFile parseProblem "input.txt"
    case s of
        Left err -> print err
        Right p  -> do
            let c = buildConstraintTensor (problemIDs p)
                constraints = sortConstraints c $ problemConstraints p
                q = map (countMatch unknownProbability . concatMap (c!) . sumQueries) constraints
            let c' = resolveConstraints constraints c
            let res = c' ! (sortedQuery . problemSpec) p
            if length res == 1
                then putStrLn $ show ((fromJust . snd . head) res) ++ "%"
                -- the direct specification wasn't resolved, let's try the inverse
                else let res' = map snd $ concatMap (c' !) $ (inverseQueryList . sortedQuery . problemSpec) p
                     in if all isJust res'
                            then putStrLn $ show (100 - (sum $ map fromJust res')) ++ "%"
                            else putStrLn "Insufficient information"

--
-- utilities
--

idOf :: ConstraintIndex -> ID
idOf (Explicit i)   = i
idOf (Inverse i)    = i
idOf (Both i)       = i

sumQueries :: Constraint -> [SortedConstraintQuery]
sumQueries (Sum x _) = x

--
-- Constraint tensor
--

buildConstraintTensor :: [ID] -> ConstraintTensor
buildConstraintTensor []     = CTEnd Nothing
buildConstraintTensor (x:xs) = CTCons x (buildConstraintTensor xs) (buildConstraintTensor xs)

sortedQuery :: ConstraintQuery -> SortedConstraintQuery
sortedQuery c = SortedConstraintQuery $ sortBy (comparing extractID) c
    where extractID (Explicit i) = i
          extractID (Inverse i) = i

(!) :: ConstraintTensor -> SortedConstraintQuery -> [QueryResult]
(!) (CTEnd x) (SortedConstraintQuery []) = [(SortedConstraintQuery [], x)]
(!) (CTCons i l r) o@(SortedConstraintQuery (x:xs))
    | i == i'   = case x of
                      Explicit _ -> subQuery l Explicit nq
                      Inverse _  -> subQuery r Inverse nq
                      Both _     -> subQuery l Explicit nq ++ subQuery r Inverse nq
    | otherwise = subQuery l Explicit o ++ subQuery r Inverse o
        where i' = idOf x
              subQuery e s q = map (\(SortedConstraintQuery c, p) -> (SortedConstraintQuery (s i : c), p)) $ e ! q
              nq = SortedConstraintQuery xs
(!) (CTCons i l r) (SortedConstraintQuery []) = subQuery l Explicit nq ++ subQuery r Inverse nq
        where nq = SortedConstraintQuery []
              subQuery e s q = map (\(SortedConstraintQuery c, p) -> (SortedConstraintQuery (s i : c), p)) $ e ! q
(!) (CTEnd _) (SortedConstraintQuery (x:xs))  = error $ "Bottomed out ConstraintTensor but query elements remain: " ++ show (x:xs)

set :: ConstraintTensor -> SortedConstraintQuery -> Probability -> ConstraintTensor
set (CTEnd _) (SortedConstraintQuery []) p = CTEnd p
set (CTCons i l r) (SortedConstraintQuery (x : xs)) p
    | i == i'   = case x of
                      Explicit _ -> CTCons i (set l nq p) r
                      Inverse _  -> CTCons i l (set r nq p)
                      Both _ -> error $ show x ++ " encountered in set, which doesn't make any sense"
    | otherwise = error $ "Failed to match ID when setting: ConstraintTensor was " ++ i ++ " but we expected " ++ i'
        where i' = idOf x
              nq = SortedConstraintQuery xs
set (CTCons i _ _) (SortedConstraintQuery []) p = error $ "Query was too short, failed at " ++ i
set (CTEnd _) (SortedConstraintQuery (x:xs)) p  = error $ "Query was too long, last id " ++ idOf x

--
-- Constraints parser
--

parseProblem :: GenParser Char st Problem
parseProblem = do
    constraintCount <- readInteger
    ids <- parseIDs
    constraints <- count constraintCount parseConstraint
    lastConstraint <- parseConstraintTree
    let constraints' = constraints ++ map inverseConstraint constraints
        constraints'' = constraints' ++ implicitConstraints ids
    return $ Problem ids (length ids) constraints'' lastConstraint

parseIDs :: GenParser Char st [ID]
parseIDs = do
    skipMany1 space
    manyTill parseID (try eol)

parseID :: GenParser Char st ID
parseID = surround (void $ many whitespace) $ many parseAlphaNum

parseConstraint :: GenParser Char st Constraint
parseConstraint = do
    cs <- parseConstraintTree
    char ':'
    many whitespace
    prob <- readFloat
    many whitespace
    eol
    return $ Sum [sortedQuery cs] (Just $ round $ 100 * prob)

parseConstraintTree :: GenParser Char st [ConstraintIndex]
parseConstraintTree =
        try parseConstraintAnd
    <|> try parseConstraintAtom
    <?> "constraint"

parseConstraintAtom =
        try parseConstraintInverse
    <|> try parseConstraintExplicit

parseConstraintAnd = chainl1 parseConstraintAtom parseConstraintAndOp
parseConstraintAndOp = symbol '&' *> pure (++)

parseConstraintInverse = do
    char '!'
    [Explicit id] <- parseConstraintExplicit
    return [Inverse id]

parseConstraintExplicit = liftA elist (many parseAlphaNum <* many whitespace)
    where elist x = [Explicit x]

readInteger :: GenParser Char st Int
readInteger = liftM (read :: String -> Int) $ many digit

readFloat :: GenParser Char st Float
readFloat = liftM (read :: String -> Float) $
    try ( do
        head <- many digit
        char '.'
        tail <- many digit
        return $ head ++ "." ++ tail )
    <|> many digit

surround :: GenParser Char st () -> GenParser Char st a -> GenParser Char st a
surround a b = a *> b <* a

symbol :: Char -> GenParser Char st Char
symbol c = many whitespace *> char c <* many whitespace

parseAlphaNum :: GenParser Char st Char
parseAlphaNum = satisfy isAlphaNum

whitespace :: GenParser Char st ()
whitespace = void $ char ' ' <|> char '\t'

eol :: GenParser Char st ()
eol = void (
        try (string "\r\n")
    <|> try (string "\n\r")
    <|> try (string "\r")
    <|> try (string "\n")
    <?> "EOL")