2022/05/Main2.idr

128 lines
5.1 KiB
Idris

import System.File.ReadWrite
import Data.Vect
import Data.String
import Data.Fin
data Crate = C Char
Show Crate where
show (C c) = pack ['[', c, ']']
parseCrate : String -> Maybe Crate
parseCrate str = let str = unpack str in
case str of
('[' :: x :: ']' :: _) => Just (C x)
_ => Nothing
data Stacks : Nat -> Type where
Stks : { n : Nat } -> Vect n (List Crate) -> Stacks n
Show (Stacks n) where
show (Stks {n} stacks) =
let longestStack = foldl max 0 . map length $ stacks
stringStacks = map (extendFront longestStack " ") . map (map show) $ stacks
lastLine = joinBy " " . map (\x => " " ++ (show x) ++ " ") $ rangeFromTo 1 n
stackLines = map (joinBy " " . toList) . transpose $ stringStacks
in joinBy "\n" . toList $ (snoc stackLines lastLine)
where extendFront : (m : Nat) -> a -> List a -> Vect m a
extendFront m x xs =
let extended = if length xs < m
then replicate (minus m (length xs)) x ++ xs
else take m xs
in case toVect m extended of
Just vect => vect
Nothing => replicate m x
transpose : { a : _ } -> Vect b (Vect a e) -> Vect a (Vect b e)
transpose [] = replicate _ []
transpose (x :: xs) = let xsTrans = transpose xs in
zipWith (::) x xsTrans
parseStacks : {n : Nat} -> List String -> Stacks n
parseStacks [] = Stks $ replicate _ []
-- Special logic for skipping the last line, since that's just the stack numbers
parseStacks (x :: []) = Stks $ replicate _ []
parseStacks {n} (x :: xs) =
let (Stks bottom) = parseStacks xs
row : Vect n (Maybe Crate)
row = map parseCrate . map pack . chunk 4 [] . unpack $ x
in Stks $ zipWith parseStacksHelper row bottom
where chunk : {n : Nat} -> (m: Nat) -> List a -> List a -> Vect n (List a)
chunk {n = 0} m defaultValue xs = []
chunk {n = (S k)} m defaultValue [] = replicate _ defaultValue
chunk {n = (S k)} m defaultValue xs =
let (head, tail) = splitAt m xs in
head :: chunk {n = k} m defaultValue tail
parseStacksHelper : Maybe Crate -> List Crate -> List Crate
parseStacksHelper (Just crate) xs = crate :: xs
parseStacksHelper Nothing xs = xs
-- Get the tops of each stack
tops : Stacks n -> String
tops (Stks xss) = pack . map (\(C c) => c) . catMaybes . toList . map head' $ xss
data Command : Nat -> Type where
Cmd : {n : Nat} -> Nat -> Fin n -> Fin n -> Command n
Show (Command n) where
show (Cmd x y z) = (show x) ++ " " ++ (show ((finToNat y) + 1)) ++ "->" ++ (show ((finToNat z) + 1))
parseCommand : {n : Nat} -> String -> Maybe (Command n)
parseCommand str =
let nums = toVect 3 . catMaybes . map parsePositive . words $ str
in do [x, y, z] <- nums
from <- natToFin (minus y 1) n
to <- natToFin (minus z 1) n
Just (Cmd x from to)
-- Apply a command to a collection of stacks
applyCommand : {n : Nat} -> Bool -> Command n -> Stacks n -> Stacks n
applyCommand multiple (Cmd count from to) stacks =
let (crates, removed) = extractCrates count from stacks
in addCrates (if multiple then crates else reverse crates) to removed
where extractCrates : Nat -> Fin n -> Stacks n -> ((List Crate), Stacks n)
extractCrates count from (Stks stacks) =
let stack = index from stacks
(head, tail) = splitAt count stack
in (head, Stks $ replaceAt from tail stacks)
addCrates : List Crate -> Fin n -> Stacks n -> Stacks n
addCrates crates to (Stks stacks) =
let stack = index to stacks
in Stks $ replaceAt to (crates ++ stack) stacks
-- Apply a list of commands to a stacks
applyCommands : {n : Nat} -> Bool -> List (Command n) -> Stacks n -> Stacks n
applyCommands multiple commands stacks = foldl (\stack, command => applyCommand multiple command stack) stacks commands
parseProblem : String -> Maybe (n ** (Stacks n, List (Command n)))
parseProblem content =
case last' . takeWhile(/= "") . lines $ content of
Just columnsLine =>
let size = foldl max 0 . catMaybes . map parsePositive . words $ columnsLine
stacks = parseStacks {n = size} . takeWhile (/= "") . lines $ content
commands = catMaybes . map (parseCommand {n = size}) . drop 1 . dropWhile (/= "") . lines $ content
in Just (size ** (stacks, commands))
Nothing => Nothing
doPart : {n : Nat} -> Bool -> Stacks n -> List (Command n) -> String
doPart multiple stacks commands = tops $ applyCommands multiple commands stacks
main : IO ()
main =
do file <- readFile "input"
case file of
Right content =>
case parseProblem content of
Just (size ** (stacks, commands)) =>
let part1 = doPart False stacks commands
part2 = doPart True stacks commands
in do putStrLn "Input:"
printLn stacks
putStr "\n\nPart 1: "
putStrLn part1
putStr "Part 2: "
putStrLn part2
Nothing => putStrLn "Failed to parse problem"
Left err => printLn err