advent05/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2022/12/04/advent-of-code-2022-day-4/
   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 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 part1 :: Wharf -> [Move] -> String
  30 part1 wharf moves = showTops $ applyMoves1 wharf moves
  31
  32 part2 :: Wharf -> [Move] -> String
  33 part2 wharf moves = showTops $ applyMoves2 wharf moves
  34
  35 showTops :: Wharf -> String
  36 showTops wharf = fmap extractName $ fmap (head . snd) $ M.toAscList wharf
  37
  38 extractName :: Crate -> Char
  39 extractName (Crate c) = c
  40
  41 makeWharf :: [[Maybe Crate]] -> [Int] -> Wharf
  42 makeWharf wharfLines colNames = M.fromList $ zip colNames wharfCols
  43   where wharfCols = fmap catMaybes $ transpose wharfLines
  44
  45 applyMoves1 :: Wharf -> [Move] -> Wharf
  46 applyMoves1 wharf moves = foldl' applyMove1 wharf moves
  47
  48 applyMove1 :: Wharf -> Move -> Wharf
  49 applyMove1 wharf m@(Move n _ _) = foldl' makeMove1 wharf (replicate n m)
  50
  51 makeMove1 :: Wharf -> Move -> Wharf
  52 makeMove1 wharf (Move _ from to) = M.insert from origin
  53                                  $ M.insert to destination wharf
  54   where (c:origin) = wharf!from
  55         destination = c:(fromMaybe [] $ wharf!?to)
  56
  57 applyMoves2 :: Wharf -> [Move] -> Wharf
  58 applyMoves2 wharf moves = foldl' applyMove2 wharf moves
  59
  60 applyMove2 :: Wharf -> Move -> Wharf
  61 applyMove2 wharf (Move n from to) = M.insert from origin'
  62                                   $ M.insert to destination wharf
  63   where origin = wharf!from
  64         moving = take n origin
  65         origin' = drop n origin
  66         destination = moving ++ (fromMaybe [] $ wharf!?to)
  67
  68
  69 -- Parse the input file
  70
  71 problemP :: Parser (([[Maybe Crate]], [Int]), [Move])
  72 wharfP :: Parser ([[Maybe Crate]], [Int])
  73 wharfLineP :: Parser [Maybe Crate]
  74 wharfCellP, blankP, crateP :: Parser (Maybe Crate)
  75 stackLabelsP :: Parser [Int]
  76 movesP :: Parser [Move]
  77 moveP :: Parser Move
  78
  79 -- problemP = (,) <$> wharfP <* endOfLine <* endOfLine <*> movesP
  80 problemP = (,) <$> wharfP <*> movesP
  81
  82 wharfP = (,) <$> (wharfLineP `sepBy` endOfLine) <*> stackLabelsP
  83
  84 wharfLineP = wharfCellP `sepBy1` (char ' ')
  85
  86 wharfCellP = crateP <|> blankP
  87 blankP = Nothing <$ (count 3 space)
  88 crateP = (Just . Crate) <$> ("[" *> letter) <* "]"
  89
  90 -- stackLabelsP = ((many1 space) *> (decimal `sepBy` (many1 space))) <* (takeWhile1 isHorizontalSpace)
  91 stackLabelsP = (many1 space)
  92             *> (decimal `sepBy` (many1 space))
  93             <* (many1 space)
  94
  95 movesP = moveP `sepBy` endOfLine
  96 moveP = Move <$> ("move " *> decimal)
  97              <*> (" from " *> decimal)
  98              <*> (" to " *> decimal)
  99
 100
 101 successfulParse :: Text -> (([[Maybe Crate]], [Int]), [Move])
 102 successfulParse input =
 103   case parseOnly problemP input of
 104     Left  _err -> (([], []), []) -- TIO.putStr $ T.pack $ parseErrorPretty err
 105     Right problem -> problem