advent05/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/04/advent-of-code-2022-day-4/
   2
   3 import System.Environment
   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 Data.Maybe
  10 import qualified Data.IntMap.Strict as M
  11 import Data.IntMap.Strict ((!), (!?))
  12
  13 data Crate = Crate Char deriving (Show, Eq)
  14 type Wharf = M.IntMap [Crate]
  15
  16 data Move = Move Int Int Int -- quantity, from, to
  17   deriving (Show, Eq)
  18
  19
  20 main :: IO ()
  21 main =
  22   do  dataFileName <- getDataFileName
  23       text <- TIO.readFile dataFileName
  24       let ((wharfLines, colNames), moves) = successfulParse text
  25       let wharf = makeWharf wharfLines colNames
  26       putStrLn $ part1 wharf moves
  27       putStrLn $ part2 wharf moves
  28
  29
  30 getDataFileName :: IO String
  31 getDataFileName =
  32   do args <- getArgs
  33      progName <- getProgName
  34      let baseDataName =  if null args
  35                          then progName
  36                          else head args
  37      let dataFileName = "data/" ++ baseDataName ++ ".txt"
  38      return dataFileName
  39
  40
  41 part1 :: Wharf -> [Move] -> String
  42 part1 wharf moves = showTops $ applyMoves1 wharf moves
  43
  44 part2 :: Wharf -> [Move] -> String
  45 part2 wharf moves = showTops $ applyMoves2 wharf moves
  46
  47 showTops :: Wharf -> String
  48 showTops wharf = fmap extractName $ fmap (head . snd) $ M.toAscList wharf
  49
  50 extractName :: Crate -> Char
  51 extractName (Crate c) = c
  52
  53 makeWharf :: [[Maybe Crate]] -> [Int] -> Wharf
  54 makeWharf wharfLines colNames = M.fromList $ zip colNames wharfCols
  55   where wharfCols = fmap (fmap fromJust)
  56                   $ fmap (dropWhile isNothing)
  57                   $ transpose wharfLines
  58
  59 applyMoves1 :: Wharf -> [Move] -> Wharf
  60 applyMoves1 wharf moves = foldl' applyMove1 wharf moves
  61
  62 applyMove1 :: Wharf -> Move -> Wharf
  63 applyMove1 wharf m@(Move n _ _) = foldl' makeMove1 wharf (replicate n m)
  64
  65 makeMove1 :: Wharf -> Move -> Wharf
  66 makeMove1 wharf (Move _ from to) = M.insert from origin
  67                                  $ M.insert to destination wharf
  68   where (c:origin) = wharf!from
  69         destination = c:(fromMaybe [] $ wharf!?to)
  70
  71 applyMoves2 :: Wharf -> [Move] -> Wharf
  72 applyMoves2 wharf moves = foldl' applyMove2 wharf moves
  73
  74 applyMove2 :: Wharf -> Move -> Wharf
  75 applyMove2 wharf (Move n from to) = M.insert from origin'
  76                                   $ M.insert to destination wharf
  77   where origin = wharf!from
  78         moving = take n origin
  79         origin' = drop n origin
  80         destination = moving ++ (fromMaybe [] $ wharf!?to)
  81
  82
  83 -- Parse the input file
  84
  85 problemP :: Parser (([[Maybe Crate]], [Int]), [Move])
  86 wharfP :: Parser ([[Maybe Crate]], [Int])
  87 wharfLineP :: Parser [Maybe Crate]
  88 wharfCellP, blankP, crateP :: Parser (Maybe Crate)
  89 stackLabelsP :: Parser [Int]
  90 movesP :: Parser [Move]
  91 moveP :: Parser Move
  92
  93 -- problemP = (,) <$> wharfP <* endOfLine <* endOfLine <*> movesP
  94 problemP = (,) <$> wharfP <*> movesP
  95
  96 wharfP = (,) <$> (wharfLineP `sepBy` endOfLine) <*> stackLabelsP
  97
  98 wharfLineP = wharfCellP `sepBy1` (char ' ')
  99
 100 wharfCellP = crateP <|> blankP
 101 blankP = Nothing <$ (count 3 space)
 102 crateP = (Just . Crate) <$> ("[" *> letter) <* "]"
 103
 104 -- stackLabelsP = ((many1 space) *> (decimal `sepBy` (many1 space))) <* (takeWhile1 isHorizontalSpace)
 105 stackLabelsP = (many1 space)
 106             *> (decimal `sepBy` (many1 space))
 107             <* (many1 space)
 108
 109 movesP = moveP `sepBy` endOfLine
 110 moveP = Move <$> ("move " *> decimal)
 111              <*> (" from " *> decimal)
 112              <*> (" to " *> decimal)
 113
 114
 115 successfulParse :: Text -> (([[Maybe Crate]], [Int]), [Move])
 116 successfulParse input =
 117   case parseOnly problemP input of
 118     Left  _err -> (([], []), []) -- TIO.putStr $ T.pack $ parseErrorPretty err
 119     Right problem -> problem