src/advent08/advent08-foldable.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 -- import Data.List
   4
   5 import Data.Monoid
   6 import Data.Text (Text)
   7 import qualified Data.Text.IO as TIO
   8
   9 import Data.Void (Void)
  10
  11 import Text.Megaparsec
  12 import Text.Megaparsec.Char
  13 import qualified Text.Megaparsec.Char.Lexer as L
  14 import qualified Control.Applicative as CA
  15
  16
  17 data Tree a = Node [Tree a] a deriving (Eq, Show)
  18
  19 instance Foldable Tree where
  20     foldMap f (Node (t: ts) m) = f m <> foldMap (foldMap f) ts
  21
  22
  23 main :: IO ()
  24 main = do
  25         text <- TIO.readFile "data/advent08-small.txt"
  26         let treeSpec = successfulParse text
  27         let (tree, _) = parseTree treeSpec
  28         print $ foldMap id tree
  29         print $ part1 tree
  30         print $ part2 tree
  31
  32 -- part1 = foldMap sum
  33 part1 = metadataOfTree
  34
  35 part2 = valueOfTree
  36
  37
  38 parseTree (c:m:spec) = (Node children metadata, remainder')
  39     where (children, remainder) = parseManyTree c spec
  40           metadata = take m remainder
  41           remainder' = drop m remainder
  42
  43 parseManyTree n spec
  44     | n == 0 = ([], spec)
  45     | otherwise = ((tree:otherTrees), remainder')
  46     where (tree, remainder) = parseTree spec
  47           (otherTrees, remainder') = parseManyTree (n-1) remainder
  48
  49 metadataOfTree (Node trees metadata) = metadata ++ (concatMap metadataOfTree trees)
  50
  51 valueOfTree (Node trees metadata)
  52     | null trees = sum metadata
  53     | otherwise = sum selectedValues
  54     where childValues = map valueOfTree trees
  55           selectedValues = map (\v -> childValues!!(v-1)) $ filter (<= (length trees)) metadata
  56
  57
  58 -- Parse the input file
  59
  60 type Parser = Parsec Void Text
  61
  62 sc :: Parser ()
  63 sc = L.space (skipSome spaceChar) CA.empty CA.empty
  64 -- sc = L.space (skipSome (char ' ')) CA.empty CA.empty
  65
  66 lexeme  = L.lexeme sc
  67 integer = lexeme L.decimal
  68
  69 treeFileP = many integer
  70
  71 successfulParse :: Text -> [Int]
  72 successfulParse input =
  73         case parse treeFileP "input" input of
  74                 Left  _error -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
  75                 Right treeSpec -> treeSpec