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