src/advent07/advent07.hs

   1 import Text.Parsec
   2 import Text.ParserCombinators.Parsec.Number
   3 import Data.List (sort, group)
   4 import qualified Data.Set as S
   5
   6 data Program = Program String Int [String]
   7                 deriving (Show, Eq)
   8
   9 name (Program n _ _) = n
  10 weight (Program _ w _) = w
  11 supports (Program _ _ s) = s
  12
  13 data Tree = Tree Program [Tree] Int deriving (Show, Eq)
  14 root (Tree p _ _) = p
  15 branches (Tree _ b _) = b
  16 tWeight (Tree _ _ w) = w
  17
  18
  19
  20 main :: IO ()
  21 main = do
  22         text <- readFile "data/advent07.txt"
  23         let progs = successfulParse $ parseFile text
  24         print $ part1 progs
  25         print $ part2 progs
  26
  27
  28 part1 :: [Program] -> String
  29 part1 progs = head $ S.elems $ S.difference pr su
  30     where su = supported progs
  31           pr = allPrograms progs
  32
  33
  34 part2 programs = (weight $ root problem) - wrongWeight + rightWeight
  35     where tree = mkTree (findByName (part1 programs) programs) programs
  36           problem = problemTree tree
  37           pt = problemParent problem tree
  38           wrongWeight = problemWeight pt
  39           rightWeight = notProblemWeight pt
  40
  41
  42 allPrograms :: [Program] -> S.Set String
  43 allPrograms = S.fromList . map name
  44
  45 supported :: [Program] -> S.Set String
  46 supported = S.unions . map (S.fromList . supports)
  47
  48
  49 -- leaves :: [Program] -> [Program]
  50 -- leaves = filter (null . supports)
  51
  52
  53 mkTree :: Program -> [Program] -> Tree
  54 mkTree program programs = Tree program subTrees (weight program + w)
  55     where subPrograms = map (\n -> findByName n programs) $ supports program
  56           subTrees = map (\r -> mkTree r programs) subPrograms
  57           w = sum $ map tWeight subTrees
  58
  59 findByName :: String -> [Program] -> Program
  60 findByName n programs = head $ filter (\p -> n == (name p)) programs
  61
  62
  63
  64 balanced :: Tree -> Bool
  65 balanced t = (S.size $ S.fromList $ map tWeight $ branches t) <= 1
  66
  67
  68 problemTree :: Tree -> Tree
  69 problemTree t
  70     | balanced t = t
  71     | otherwise = problemTree problemSubtree
  72         where subtreeWeights = map tWeight $ branches t
  73               weightGroups = group $ sort subtreeWeights
  74               pWeight = head $ head $ filter (\g -> length g == 1) weightGroups
  75               problemSubtree = head $ filter (\s -> tWeight s == pWeight) (branches t)
  76
  77
  78 problemParent :: Tree -> Tree -> Tree
  79 problemParent problem tree = head $ problemParent' problem tree
  80
  81 problemParent' :: Tree -> Tree -> [Tree]
  82 problemParent' problem tree
  83     | problem `elem` (branches tree) = [tree]
  84     | null $ branches tree = []
  85     | otherwise = concatMap (problemParent' problem) $ branches tree
  86
  87
  88 problemWeight :: Tree -> Int
  89 problemWeight tree = head $ head $ filter (\g -> 1 == length g) $ group $ sort $ map tWeight $ branches tree
  90
  91 notProblemWeight :: Tree -> Int
  92 notProblemWeight tree = head $ head $ filter (\g -> 1 /= length g) $ group $ sort $ map tWeight $ branches tree
  93
  94
  95
  96 onlySpaces = many (oneOf " \t")
  97 parens = between (string "(") (string ")")
  98 symP = many lower
  99 commaSep sym = sym `sepBy` (onlySpaces *> string "," *> onlySpaces)
 100
 101 mFile = mLine `sepBy` newline
 102 mLine = Program <$> symP <*> (onlySpaces *> (parens int)) <*> supportsP
 103 supportsP = (onlySpaces *> (string "->") *> onlySpaces *> (commaSep symP)) <|> (pure [])
 104
 105 parseFile :: String -> Either ParseError [Program]
 106 parseFile input = parse mFile "(unknown)" input
 107
 108 -- parseLine :: String -> Either ParseError Program
 109 -- parseLine input = parse mLine "(unknown)" input
 110
 111 successfulParse :: Either ParseError [a] -> [a]
 112 successfulParse (Left _) = []
 113 successfulParse (Right a) = a
 114
 115
 116
 117 -- sampleT = "pbga (66)\nxhth (57)\nebii (61)\nhavc (66)\nktlj (57)\nfwft (72) -> ktlj, cntj, xhth\nqoyq (66)\npadx (45) -> pbga, havc, qoyq\ntknk (41) -> ugml, padx, fwft\njptl (61)\nugml (68) -> gyxo, ebii, jptl\ngyxo (61)\ncntj (57)"
 118 -- sample = successfulParse $ parseFile sampleT
 119
 120 -- sampleLeaves = leaves sample
 121 -- sampleBranch = sample \\ sampleLeaves