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 catMaybes $ transpose wharfLines
  56
  57 applyMoves1 :: Wharf -> [Move] -> Wharf
  58 applyMoves1 wharf moves = foldl' applyMove1 wharf moves
  59
  60 applyMove1 :: Wharf -> Move -> Wharf
  61 applyMove1 wharf m@(Move n _ _) = foldl' makeMove1 wharf (replicate n m)
  62
  63 makeMove1 :: Wharf -> Move -> Wharf
  64 makeMove1 wharf (Move _ from to) = M.insert from origin
  65                                  $ M.insert to destination wharf
  66   where (c:origin) = wharf!from
  67         destination = c:(fromMaybe [] $ wharf!?to)
  68
  69 applyMoves2 :: Wharf -> [Move] -> Wharf
  70 applyMoves2 wharf moves = foldl' applyMove2 wharf moves
  71
  72 applyMove2 :: Wharf -> Move -> Wharf
  73 applyMove2 wharf (Move n from to) = M.insert from origin'
  74                                   $ M.insert to destination wharf
  75   where origin = wharf!from
  76         moving = take n origin
  77         origin' = drop n origin
  78         destination = moving ++ (fromMaybe [] $ wharf!?to)
  79
  80
  81 -- Parse the input file
  82
  83 problemP :: Parser (([[Maybe Crate]], [Int]), [Move])
  84 wharfP :: Parser ([[Maybe Crate]], [Int])
  85 wharfLineP :: Parser [Maybe Crate]
  86 wharfCellP, blankP, crateP :: Parser (Maybe Crate)
  87 stackLabelsP :: Parser [Int]
  88 movesP :: Parser [Move]
  89 moveP :: Parser Move
  90
  91 -- problemP = (,) <$> wharfP <* endOfLine <* endOfLine <*> movesP
  92 problemP = (,) <$> wharfP <*> movesP
  93
  94 wharfP = (,) <$> (wharfLineP `sepBy` endOfLine) <*> stackLabelsP
  95
  96 wharfLineP = wharfCellP `sepBy1` (char ' ')
  97
  98 wharfCellP = crateP <|> blankP
  99 blankP = Nothing <$ (count 3 space)
 100 crateP = (Just . Crate) <$> ("[" *> letter) <* "]"
 101
 102 -- stackLabelsP = ((many1 space) *> (decimal `sepBy` (many1 space))) <* (takeWhile1 isHorizontalSpace)
 103 stackLabelsP = (many1 space)
 104             *> (decimal `sepBy` (many1 space))
 105             <* (many1 space)
 106
 107 movesP = moveP `sepBy` endOfLine
 108 moveP = Move <$> ("move " *> decimal)
 109              <*> (" from " *> decimal)
 110              <*> (" to " *> decimal)
 111
 112
 113 successfulParse :: Text -> (([[Maybe Crate]], [Int]), [Move])
 114 successfulParse input =
 115   case parseOnly problemP input of
 116     Left  _err -> (([], []), []) -- TIO.putStr $ T.pack $ parseErrorPretty err
 117     Right problem -> problem