src/advent08/advent08-treelibrary.hs

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