advent05/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/05/advent-of-code-2023-day-05/
   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.List
   9 import qualified Data.Map.Strict as M
  10 import Data.Map.Strict ((!))
  11 import Data.List.Split (chunksOf)
  12
  13 type Almanac = M.Map String AMap
  14 data AMap = AMap String [Rule] deriving (Eq, Show)
  15 data Rule = Rule Interval Int deriving (Eq, Ord, Show)
  16 data Requirement = Requirement String [Interval] deriving (Eq, Show)
  17
  18 data Interval = Iv Int Int deriving (Eq, Ord, Show) -- inclusive, closed at both ends
  19
  20
  21 main :: IO ()
  22 main =
  23   do  dataFileName <- getDataFileName
  24       text <- TIO.readFile dataFileName
  25       let (seeds, almanac) = successfulParse text
  26       print $ part1 almanac seeds
  27       print $ part2 almanac seeds
  28
  29 part1, part2 :: Almanac -> [Int] -> Int
  30 part1 almanac seeds = lowestLocation almanac seeds'
  31   where seeds' = tidyIntervals $ singletonRanges seeds
  32
  33 part2 almanac seeds = lowestLocation almanac seeds'
  34   where seeds' = tidyIntervals $ expandRanges seeds
  35
  36 lowestLocation :: Almanac -> [Interval] -> Int
  37 lowestLocation almanac seeds = l
  38   where Requirement _ locations = followRequirements almanac $ Requirement "seed" seeds
  39         (Iv l _) = head locations
  40
  41 followRequirements :: Almanac -> Requirement -> Requirement
  42 followRequirements _ req@(Requirement "location" vals) = req
  43 followRequirements almanac (Requirement name vals) =
  44   followRequirements almanac newReq
  45   where aMap = almanac ! name
  46         newReq = useAMap aMap vals
  47
  48
  49 useRule :: Rule -> Interval -> ([Interval], [Interval], [Rule])
  50 useRule (Rule (Iv rl rh) d) (Iv xl xh) = (newResults, newVals, newRules)
  51   where newResults =
  52           filter legalInterval [ Iv (min xl rl) (min xh (rl - 1)) -- input below rule
  53                                , Iv ((max xl rl) + d) ((min xh rh) + d)] -- input within rule
  54         newVals = filter legalInterval [Iv (max xl (rh + 1)) (max xh rh)] -- input above rule
  55         newRules = filter legalRule [Rule (Iv (max (xh + 1) rl) (max xh rh)) d] -- rule above input
  56
  57
  58 useRules :: [Rule] -> [Interval] -> [Interval]
  59 useRules [] vals = vals
  60 useRules _ [] = []
  61 useRules (r@(Rule (Iv rl rh) _):rs) (v@(Iv xl xh):vs)
  62   | rh < xl = useRules rs (v:vs)
  63   | xh < rl = v : useRules (r:rs) vs
  64   | otherwise = newResults ++ (useRules (newRules ++ rs) (newVals ++ vs))
  65   where (newResults, newVals, newRules) = useRule r v
  66
  67
  68 legalInterval :: Interval -> Bool
  69 legalInterval (Iv l h) = l <= h
  70
  71 legalRule :: Rule -> Bool
  72 legalRule (Rule iv _) = legalInterval iv
  73
  74 useAMap :: AMap -> [Interval] -> Requirement
  75 useAMap (AMap d rs) xs = Requirement d $ tidyIntervals $ useRules rs xs
  76
  77
  78 singletonRanges :: [Int] -> [Interval]
  79 singletonRanges = fmap (\x -> Iv x x)
  80
  81 expandRanges :: [Int] -> [Interval]
  82 expandRanges seeds = fmap expandRange ranges
  83   where ranges = chunksOf 2 seeds
  84         expandRange [s, l] = Iv s (s + l - 1)
  85
  86
  87 tidyIntervals :: [Interval] -> [Interval]
  88 tidyIntervals ivs0 = tidyIntervalsS $ sort ivs0
  89
  90 tidyIntervalsS :: [Interval] -> [Interval]
  91 tidyIntervalsS [] = []
  92 tidyIntervalsS [x] = [x]
  93 tidyIntervalsS (x:y:xs)
  94   | x `allBelow` y = x : tidyIntervalsS (y:xs)
  95   | otherwise = tidyIntervalsS ((x `merge` y) : xs)
  96
  97 allBelow :: Interval -> Interval -> Bool
  98 allBelow (Iv _ x2) (Iv y1 _) = (x2 + 1) < y1
  99
 100 merge :: Interval -> Interval -> Interval
 101 merge (Iv x1 x2) (Iv y1 y2) = Iv (min x1 y1) (max x2 y2)
 102
 103
 104
 105 -- Parse the input file
 106
 107 problemP :: Parser ([Int], Almanac)
 108 seedsP :: Parser [Int]
 109 almanacP :: Parser Almanac
 110 aMapP :: Parser (String, AMap)
 111 aMapHeaderP :: Parser (String, String)
 112 rulesP :: Parser [Rule]
 113 ruleP :: Parser Rule
 114 numbersP :: Parser [Int]
 115 nameP :: Parser String
 116 blankLineP :: Parser ()
 117
 118
 119 problemP = (,) <$> (seedsP <* blankLineP) <*> almanacP
 120
 121 seedsP = "seeds: " *> numbersP
 122
 123 almanacP = M.fromList <$> (aMapP `sepBy` blankLineP)
 124
 125 aMapP = aMapify <$> aMapHeaderP <*> rulesP
 126 aMapHeaderP = (,) <$> (nameP <* "-to-") <*> (nameP <* " map:" <* endOfLine)
 127
 128 rulesP = ruleP `sepBy` endOfLine
 129 ruleP = ruleify <$> (decimal <* space) <*> (decimal <* space) <*> decimal
 130
 131 numbersP = decimal `sepBy` skipSpace
 132 nameP = many1 letter
 133
 134 blankLineP = endOfLine *> endOfLine
 135
 136 aMapify :: (String, String) -> [Rule] -> (String, AMap)
 137 aMapify (s, d) rs = (s, AMap d (sort rs))
 138
 139 ruleify :: Int -> Int -> Int -> Rule
 140 ruleify d s l = Rule (Iv s (s + l - 1)) (d - s)
 141
 142 successfulParse :: Text -> ([Int], Almanac)
 143 successfulParse input =
 144   case parseOnly problemP input of
 145     Left  _err -> ([], M.empty) -- TIO.putStr $ T.pack $ parseErrorPretty err
 146     Right matches -> matches