Advent of Code - 2021 - 23

Haskell

import qualified Data.PQueue.Prio.Min as P
import qualified Data.Array  as A
import qualified Data.IntMap as M
import qualified Data.Set    as S
import Data.Maybe (isJust, isNothing)
import Data.Maybe (catMaybes, fromJust)

data Amphipod = A | B | C | D deriving (Eq, Ord, Show)
-- #01x2x3x4x56#
type Hallway  = A.Array Int (Maybe Amphipod)
data Room     = Z
              | F Amphipod
              | L Amphipod
              | FF Amphipod Amphipod
              | LL Amphipod Amphipod
              | FFF Amphipod Amphipod Amphipod
              | LLL Amphipod Amphipod Amphipod
              | FFFF Amphipod Amphipod Amphipod Amphipod
              | LLLL Amphipod Amphipod Amphipod Amphipod
              deriving (Eq, Ord, Show)
type Rooms    = A.Array Int Room
data Burrow   = Burrow Hallway Rooms deriving (Ord, Eq, Show)
type RoomId     = Int
type HallwayPos = Int
type Cost = Int

cost A = 1
cost B = 10
cost C = 100
cost D = 1000

room (a,b,c,d) = FFFF a b c d

burrow a b c d = Burrow (A.listArray (0,10) $ repeat Nothing)
                        (A.listArray (0,3) $ map room [a,b,c,d])
         
-- I was curious about other peoples so I took a bunch of inputs and
-- measured them
input   = burrow (D,D,D,B) (C,B,C,B) (A,A,B,C) (A,C,A,D) -- 1m01s
input'0 = burrow (D,D,D,B) (A,B,C,A) (D,A,B,B) (C,C,A,C) --   35s
input'1 = burrow (C,D,D,D) (C,B,C,B) (A,A,B,D) (B,C,A,A) -- N/A (probably has an error)
test    = burrow (A,D,D,B) (D,B,C,C) (C,A,B,B) (A,C,A,D) -- 1m08s

pushRoom e (Z        ) = Just (L    e      )
pushRoom e (F   a    ) = Just (LL   a e    )
pushRoom e (L   a    ) = Just (LL   a e    )
pushRoom e (FF  a b  ) = Just (LLL  a b e  )
pushRoom e (LL  a b  ) = Just (LLL  a b e  )
pushRoom e (FFF a b c) = Just (LLLL a b c e)
pushRoom e (LLL a b c) = Just (LLLL a b c e)
pushRoom e _           = Nothing

popRoom (FFFF a b c d) = Just (d, FFF a b c)
popRoom (FFF  a b c  ) = Just (c, FF  a b  )
popRoom (FF   a b    ) = Just (b, F   a    )
popRoom (F    a      ) = Just (a, Z        )
popRoom _              = Nothing

putHallway ri e hi hw | any isJust c = Nothing
                      | otherwise    = Just $ hw A.// [(hi, Just e)]
  where b x y = [min x y .. max x y]
        c = [hw A.! i | i <- b (roomPos ri) hi]

takeHallway ri hi hw | any isJust c        = Nothing
                     | Just e <- hw A.! hi = Just (e, hw A.// [(hi, Nothing)])
                     | otherwise           = Nothing
  where b x y = [min x y .. max x y]
        c = [hw A.! i | i <- b (roomPos ri) hi, i /= hi]

roomPos 0 = 2
roomPos 1 = 4
roomPos 2 = 6
roomPos 3 = 8

roomPositions = [2,4,6,8]
hallwayPositions = [0,1,3,5,7,9,10]

distance ri hi = abs $ roomPos ri - hi

moveToHallway :: RoomId -> HallwayPos -> Burrow -> Maybe (Burrow, Cost)
moveToHallway ri hi (Burrow hw rms) = do
  (e,rm) <- popRoom $ rms A.! ri
  hw     <- putHallway ri e hi hw
  let c = (distance ri hi + roomSteps rm) * cost e
  pure (Burrow hw (rms A.// [(ri, rm)]), c)
  where roomSteps (Z        ) = 4
        roomSteps (F   _    ) = 3
        roomSteps (L   _    ) = 3
        roomSteps (FF  _ _  ) = 2
        roomSteps (LL  _ _  ) = 2
        roomSteps (FFF _ _ _) = 1
        roomSteps (LLL _ _ _) = 1

moveToRoom :: HallwayPos -> RoomId -> Burrow -> Maybe (Burrow, Cost)
moveToRoom hi ri (Burrow hw rms) = do
  (e,hw) <- takeHallway ri hi hw
  rm     <- pushRoom e $ rms A.! ri
  let c = (distance ri hi + roomSteps rm) * cost e
  pure (Burrow hw (rms A.// [(ri, rm)]), c)
  where roomSteps (L    _      ) = 4
        roomSteps (LL   _ _    ) = 3
        roomSteps (LLL  _ _ _  ) = 2
        roomSteps (LLLL _ _ _ _) = 1

solved (Burrow _ rm) = f A 0 && f B 1 && f C 2 && f D 3
  where f a r = rm A.! r == LLLL a a a a

unsolvable (Burrow _ rm) = f A 0 || f B 1 || f C 2 || f D 3
  where f a r = g a $ rm A.! r
        g e (LLLL a b c d) = a /= e || b /= e || c /= e || d /= e
        g e (LLL  a b c  ) = a /= e || b /= e || c /= e
        g e (LL   a b    ) = a /= e || b /= e
        g e (L    a      ) = a /= e
        g e _              = False

solve :: Burrow -> Cost
solve b = loop S.empty $ P.singleton 0 b
  where
    loop :: S.Set Burrow -> P.MinPQueue Cost Burrow -> Cost
    loop s m | solved b  = c
             | otherwise = loop s' m''
      where ((c, b), m') = fromJust $ P.minViewWithKey m
            lh  = [moveToHallway r h b | r <- [0..3], h <- hallwayPositions]
            lr  = [moveToRoom    h r b | r <- [0..3], h <- hallwayPositions]
            l   = h (lr ++ lh)
            h = filter (\(b,c) -> S.notMember b s)
              . filter (not . unsolvable . fst)
              . catMaybes
            m'' = foldr (uncurry P.insert) m' $ map (\(b,c') -> (c + c', b)) l
            s' = S.insert b s

main = print $ solve input