128 lines
5.1 KiB
Idris
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
|