work in progress
[advent-of-code-22.git] / advent07 / Main.hs
1 -- Writeup at https://work.njae.me.uk/2022/12/07/advent-of-code-2022-day-7/
2
3 import AoC
4 import Data.Text (Text)
5 import qualified Data.Text.IO as TIO
6 import Data.Attoparsec.Text hiding (take)
7 import Control.Applicative
8 import Data.Char
9 import Data.Maybe
10 import Data.Tree
11 import Data.Tree.Zipper hiding (tree)
12 import qualified Data.Map.Strict as M
13 -- import Data.Map.Strict ((!), (!?))
14 import Data.List (foldl', sort)
15
16 data ParsedObject = CD String
17 | LS
18 | PDirectory String
19 | PFile Int String
20 deriving Show
21
22 data Directory = Dir String (M.Map String Int)
23 deriving (Show, Eq)
24
25 data ContainedSize = CSize String Integer
26 deriving (Show, Eq)
27
28 type DTree = Tree Directory
29 type ZDTree = TreePos Full Directory
30 type STree = Tree ContainedSize
31
32 reportingThreshold, spaceAvailable, spaceRequired :: Integer
33 reportingThreshold = 100000
34 spaceAvailable = 70000000
35 spaceRequired = 30000000
36
37 main :: IO ()
38 main =
39 do dataFileName <- getDataFileName
40 text <- TIO.readFile dataFileName
41 let trace = successfulParse text
42 let tree = mkTree trace emptyTree
43 let sizedTree = transitiveSizes $ containingSizes tree
44 print $ part1 sizedTree
45 print $ part2 sizedTree
46
47 part1, part2 :: STree -> Integer
48 part1 = foldTree (\x xs -> sum (x:xs)) . fmap cancelLarge
49
50 part2 tree = spaceFreed
51 where nodes = fmap extractCSize $ flatten tree
52 spaceUsed = extractCSize $ rootLabel tree
53 spaceUnused = spaceAvailable - spaceUsed
54 spaceToFree = spaceRequired - spaceUnused
55 viableNodes = filter (>= spaceToFree) nodes
56 spaceFreed = head $ sort viableNodes
57
58 emptyTree :: DTree
59 emptyTree = Node {rootLabel = (Dir "/" M.empty), subForest = []}
60
61 mkTree :: [ParsedObject] -> DTree -> DTree
62 mkTree trace tree = toTree $ root $ makeTree trace $ fromTree tree
63
64 makeTree :: [ParsedObject] -> ZDTree -> ZDTree
65 makeTree trace tree = foldl' processCommand tree trace
66
67 processCommand :: ZDTree -> ParsedObject -> ZDTree
68 processCommand tree (CD name)
69 | name == "/" = root tree
70 | name == ".." = fromJust $ parent tree
71 | otherwise = fromJust $ childWithName name tree
72 processCommand tree LS = tree
73 processCommand tree (PFile size name) =
74 modifyLabel (\ (Dir n fs) -> Dir n (M.insert name size fs)) tree
75 processCommand tree (PDirectory name) =
76 if (isJust $ childWithName name tree)
77 then tree
78 else fromJust $ parent $ insert (Node { rootLabel = (Dir name M.empty)
79 , subForest = []
80 }) $ children tree
81
82
83 childWithName :: String -> ZDTree -> Maybe ZDTree
84 childWithName name tree = searchForChild name (firstChild tree)
85
86 searchForChild :: String -> Maybe ZDTree -> Maybe ZDTree
87 searchForChild _name Nothing = Nothing
88 searchForChild name (Just tree)
89 | name == labelName = Just tree
90 | otherwise = searchForChild name (next tree)
91 where (Dir labelName _) = label tree
92
93 containingSizes :: DTree -> STree
94 containingSizes (Node {rootLabel = (Dir name files), subForest = sf}) =
95 (Node {rootLabel = (CSize name sizeHere), subForest = sizedTrees})
96 where sizeHere = M.foldl (+) 0 $ M.map fromIntegral files
97 sizedTrees = fmap containingSizes sf
98
99 transitiveSizes :: STree -> STree
100 transitiveSizes (Node {rootLabel = (CSize name sizeHere), subForest = sf}) =
101 (Node {rootLabel = (CSize name (sizeHere + subSizes)), subForest = sizedTrees })
102 where sizedTrees = fmap transitiveSizes sf
103 subSizes = sum $ fmap (extractCSize . rootLabel) sizedTrees
104
105 extractCSize, cancelLarge :: ContainedSize -> Integer
106 extractCSize (CSize _ s) = s
107
108 cancelLarge (CSize _ s)
109 | s <= reportingThreshold = s
110 | otherwise = 0
111
112 -- Parse the input file
113
114 traceP :: Parser [ParsedObject]
115 lineP :: Parser ParsedObject
116 cdP :: Parser ParsedObject
117 lsP :: Parser ParsedObject
118 fileP :: Parser ParsedObject
119 directoryP :: Parser ParsedObject
120 letterP :: Parser Char
121 nameP :: Parser String
122
123 traceP = lineP `sepBy` endOfLine
124
125 lineP = cdP <|> lsP <|> directoryP <|> fileP
126 cdP = CD <$> ("$ cd " *> nameP)
127 lsP = LS <$ "$ ls"
128 fileP = PFile <$> (decimal <* " ") <*> nameP
129 directoryP = PDirectory <$> ("dir " *> nameP)
130
131 nameP = many1 letterP
132 letterP = satisfy (not . isSpace)
133
134 successfulParse :: Text -> [ParsedObject]
135 successfulParse input =
136 case parseOnly traceP input of
137 Left _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
138 Right commands -> commands