158 lines
6.1 KiB
Haskell
Executable File
158 lines
6.1 KiB
Haskell
Executable File
#! /usr/bin/env stack
|
|
-- stack --resolver lts-12.20 --install-ghc script
|
|
{-# LANGUAGE ViewPatterns, BangPatterns, LambdaCase, TemplateHaskell #-}
|
|
module Main where
|
|
|
|
import Control.Applicative
|
|
import Control.Lens
|
|
import Control.Monad.State
|
|
import Data.Function (on)
|
|
import Data.Maybe (fromJust)
|
|
import Data.Tuple (swap)
|
|
import Data.Void
|
|
import System.Environment (getArgs)
|
|
|
|
import qualified Data.Array as A
|
|
import qualified Data.List as L
|
|
import qualified Data.Map as M
|
|
import qualified Data.Set as S
|
|
|
|
import Debug.Trace
|
|
|
|
type CartID = Int
|
|
type Position = (Int, Int)
|
|
|
|
data Direction = DUp | DRight | DDown | DLeft
|
|
deriving (Eq,Ord,Enum,Bounded,Show)
|
|
|
|
data Bias = TurnLeft | GoStraight | TurnRight
|
|
deriving (Eq,Ord,Enum,Bounded,Show)
|
|
|
|
data Cart = Cart { _cartPosition :: Position
|
|
, _cartDirection :: Direction
|
|
, _cartBias :: Bias
|
|
} deriving (Eq,Ord,Show)
|
|
makeLenses ''Cart
|
|
|
|
type Carts = M.Map CartID Cart
|
|
|
|
data Track = TEmpty | THorizontal | TVertical
|
|
| TCurvedFW | TCurvedBW | TIntersection
|
|
deriving (Eq,Ord,Enum,Bounded,Show)
|
|
|
|
type Tracks = A.Array Position Track
|
|
|
|
data Part = Part1 | Part2 deriving (Eq,Ord,Enum,Bounded,Show)
|
|
|
|
main = getArgs >>= \case
|
|
[] -> interact $ solve Part1 . parseInput
|
|
["1"] -> interact $ solve Part1 . parseInput
|
|
["2"] -> interact $ solve Part2 . parseInput
|
|
_ -> putStrLn "Invalid arguments."
|
|
|
|
parseInput :: String -> (Tracks, Carts)
|
|
parseInput input = (tracks, carts)
|
|
where
|
|
max_x = (maximum $ map length $ lines input) - 1
|
|
max_y = (length $ lines input) - 1
|
|
tracks = A.accumArray (flip const) TEmpty ((0,0),(max_x,max_y)) $
|
|
[ ((x, y), case c of { '\\' -> TCurvedBW
|
|
; '/' -> TCurvedFW
|
|
; '-' -> THorizontal
|
|
; '<' -> THorizontal
|
|
; '>' -> THorizontal
|
|
; '|' -> TVertical
|
|
; '^' -> TVertical
|
|
; 'v' -> TVertical
|
|
; 'V' -> TVertical
|
|
; '+' -> TIntersection
|
|
})
|
|
| (y, ln) <- zip [0..] $ lines input
|
|
, (x, c) <- zip [0..] ln
|
|
, c /= ' '
|
|
]
|
|
carts = M.fromList $ zip [0..] $
|
|
[ Cart (x, y) direction TurnLeft
|
|
| (y, ln) <- zip [0..] $ lines input
|
|
, (x, c) <- zip [0..] ln
|
|
, direction <- case c of { '^' -> [DUp]
|
|
; '>' -> [DRight]
|
|
; 'v' -> [DDown]
|
|
; '<' -> [DLeft]
|
|
; _ -> []
|
|
}
|
|
]
|
|
|
|
compareOrder :: Carts -> CartID -> CartID -> Ordering
|
|
compareOrder carts = compare `on` (\i -> carts ^? at i . to fromJust . cartPosition . to swap)
|
|
|
|
solve :: Part -> (Tracks, Carts) -> String
|
|
solve part (tracks, initialCarts) = result
|
|
where
|
|
atPosition :: Position -> Track
|
|
atPosition pt
|
|
| A.inRange (A.bounds tracks) pt = tracks A.! pt
|
|
| otherwise = TEmpty
|
|
result = either id absurd $ flip evalStateT initialCarts $ forever $ do
|
|
let sortCartIDs carts = L.sortBy (compareOrder carts) (M.keys carts)
|
|
sortedCarts <- sortCartIDs <$> get
|
|
forM_ sortedCarts $ \cartID -> do
|
|
removed <- M.notMember cartID <$> get
|
|
unless removed $ do
|
|
-- move cart
|
|
direction <- use $ unsafeIx cartID . cartDirection
|
|
modifying (ix cartID . cartPosition) $ \(x,y) ->
|
|
case direction of
|
|
DUp -> (x, y-1)
|
|
DRight -> (x+1, y)
|
|
DDown -> (x, y+1)
|
|
DLeft -> (x-1, y)
|
|
-- turn cart
|
|
position <- use $ unsafeIx cartID . cartPosition
|
|
case atPosition position of
|
|
TIntersection -> do
|
|
bias <- use $ unsafeIx cartID . cartBias
|
|
ix cartID . cartDirection .= case (direction, bias) of
|
|
(DUp, TurnLeft) -> DLeft
|
|
(DUp, GoStraight) -> DUp
|
|
(DUp, TurnRight) -> DRight
|
|
(DRight, TurnLeft) -> DUp
|
|
(DRight, GoStraight) -> DRight
|
|
(DRight, TurnRight) -> DDown
|
|
(DDown, TurnLeft) -> DRight
|
|
(DDown, GoStraight) -> DDown
|
|
(DDown, TurnRight) -> DLeft
|
|
(DLeft, TurnLeft) -> DDown
|
|
(DLeft, GoStraight) -> DLeft
|
|
(DLeft, TurnRight) -> DUp
|
|
ix cartID . cartBias .= case bias of
|
|
TurnLeft -> GoStraight
|
|
GoStraight -> TurnRight
|
|
TurnRight -> TurnLeft
|
|
TCurvedFW -> do -- '/'
|
|
ix cartID . cartDirection .= case direction of
|
|
DUp -> DRight
|
|
DRight -> DUp
|
|
DDown -> DLeft
|
|
DLeft -> DDown
|
|
TCurvedBW -> do -- '\'
|
|
ix cartID . cartDirection .= case direction of
|
|
DUp -> DLeft
|
|
DRight -> DDown
|
|
DDown -> DRight
|
|
DLeft -> DUp
|
|
_ -> pure ()
|
|
-- check for collisions
|
|
here <- map fst . filter (\(i,c) -> c^.cartPosition == position) . M.assocs <$> get
|
|
when (length here > 1) $ do
|
|
when (part == Part1) $ lift $ Left $ show position
|
|
forM_ here $ modify . M.delete
|
|
when (part == Part2) $ do
|
|
remaining <- toListOf (each . cartPosition) <$> get
|
|
--traceM $ show remaining
|
|
when (null (drop 1 remaining)) $ lift $ Left $ show remaining
|
|
|
|
unsafeIx :: (Functor f, Ixed t) => Index t
|
|
-> Over (->) f t t (IxValue t) (IxValue t)
|
|
unsafeIx = unsafeSingular . ix
|