{-# Language RecordWildCards, ImportQualifiedPost #-}
module Main (main) where
import Advent (getInputArray, countBy)
import Advent.Coord (above, below, cardinal, left, right, Coord)
import Advent.Search (astar, astarOn, bfsOn, AStep(..))
import Data.Array.Unboxed ( UArray, (!), (//), assocs, elems )
import Data.Char (ord, isLower, isUpper)
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.List ( foldl' )
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Map.Strict qualified as MapStrict
import Data.Maybe (isJust)
import Data.Set (Set)
import Data.Set qualified as Set
main :: IO ()
IO ()
main =
do UArray Coord Char
world1 <- Int -> Int -> IO (UArray Coord Char)
getInputArray Int
2019 Int
18
let start :: Coord
start = [Coord] -> Coord
forall a. HasCallStack => [a] -> a
head [Coord
k | (Coord
k,Char
'@') <- UArray Coord Char -> [(Coord, Char)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Coord Char
world1]
Int -> IO ()
forall a. Show a => a -> IO ()
print (UArray Coord Char -> [Coord] -> Int
allKeys UArray Coord Char
world1 [Coord
start])
let fixups :: [(Coord, Char)]
fixups = [(Coord
c,Char
'#') | Coord
c <- Coord
start Coord -> [Coord] -> [Coord]
forall a. a -> [a] -> [a]
: Coord -> [Coord]
cardinal Coord
start]
[(Coord, Char)] -> [(Coord, Char)] -> [(Coord, Char)]
forall a. [a] -> [a] -> [a]
++ [(Coord -> Coord
f (Coord -> Coord
g Coord
start),Char
'@') | Coord -> Coord
f <- [Coord -> Coord
above, Coord -> Coord
below], Coord -> Coord
g <- [Coord -> Coord
left , Coord -> Coord
right]]
world2 :: UArray Coord Char
world2 = UArray Coord Char
world1 UArray Coord Char -> [(Coord, Char)] -> UArray Coord Char
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)] -> a i e
// [(Coord, Char)]
fixups
start2 :: [Coord]
start2 = [Coord
k | (Coord
k,Char
'@') <- UArray Coord Char -> [(Coord, Char)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Coord Char
world2]
Int -> IO ()
forall a. Show a => a -> IO ()
print (UArray Coord Char -> [Coord] -> Int
allKeys UArray Coord Char
world2 [Coord]
start2)
data Cell = Start | Gate !Int | Key !Int
deriving (Cell -> Cell -> Bool
(Cell -> Cell -> Bool) -> (Cell -> Cell -> Bool) -> Eq Cell
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Cell -> Cell -> Bool
== :: Cell -> Cell -> Bool
$c/= :: Cell -> Cell -> Bool
/= :: Cell -> Cell -> Bool
Eq, Int -> Cell -> ShowS
[Cell] -> ShowS
Cell -> String
(Int -> Cell -> ShowS)
-> (Cell -> String) -> ([Cell] -> ShowS) -> Show Cell
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Cell -> ShowS
showsPrec :: Int -> Cell -> ShowS
$cshow :: Cell -> String
show :: Cell -> String
$cshowList :: [Cell] -> ShowS
showList :: [Cell] -> ShowS
Show, Eq Cell
Eq Cell =>
(Cell -> Cell -> Ordering)
-> (Cell -> Cell -> Bool)
-> (Cell -> Cell -> Bool)
-> (Cell -> Cell -> Bool)
-> (Cell -> Cell -> Bool)
-> (Cell -> Cell -> Cell)
-> (Cell -> Cell -> Cell)
-> Ord Cell
Cell -> Cell -> Bool
Cell -> Cell -> Ordering
Cell -> Cell -> Cell
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Cell -> Cell -> Ordering
compare :: Cell -> Cell -> Ordering
$c< :: Cell -> Cell -> Bool
< :: Cell -> Cell -> Bool
$c<= :: Cell -> Cell -> Bool
<= :: Cell -> Cell -> Bool
$c> :: Cell -> Cell -> Bool
> :: Cell -> Cell -> Bool
$c>= :: Cell -> Cell -> Bool
>= :: Cell -> Cell -> Bool
$cmax :: Cell -> Cell -> Cell
max :: Cell -> Cell -> Cell
$cmin :: Cell -> Cell -> Cell
min :: Cell -> Cell -> Cell
Ord)
charToCell :: Char -> Maybe Cell
charToCell :: Char -> Maybe Cell
charToCell Char
x
| Char
'@' Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
x = Cell -> Maybe Cell
forall a. a -> Maybe a
Just Cell
Start
| Char -> Bool
isLower Char
x = Cell -> Maybe Cell
forall a. a -> Maybe a
Just (Int -> Cell
Key (Char -> Int
ord Char
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Char -> Int
ord Char
'a'))
| Char -> Bool
isUpper Char
x = Cell -> Maybe Cell
forall a. a -> Maybe a
Just (Int -> Cell
Gate (Char -> Int
ord Char
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Char -> Int
ord Char
'A'))
| Bool
otherwise = Maybe Cell
forall a. Maybe a
Nothing
extractGraph :: UArray Coord Char -> Map Coord [(Coord, Cell, Int)]
UArray Coord Char
world =
[(Coord, [(Coord, Cell, Int)])] -> Map Coord [(Coord, Cell, Int)]
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
[ (Coord
pos, UArray Coord Char -> Coord -> Cell -> [(Coord, Cell, Int)]
startSearch UArray Coord Char
world Coord
pos Cell
cell)
| (Coord
pos, Char
char) <- UArray Coord Char -> [(Coord, Char)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs (UArray Coord Char
world :: UArray Coord Char)
, Just Cell
cell <- [Char -> Maybe Cell
charToCell Char
char]
]
startSearch :: UArray Coord Char -> Coord -> Cell -> [(Coord, Cell, Int)]
startSearch :: UArray Coord Char -> Coord -> Cell -> [(Coord, Cell, Int)]
startSearch UArray Coord Char
world Coord
start Cell
startCell =
[ (Coord
here, Cell
cell, Int
n)
| (Coord
here, Just Cell
cell, Int
n) <- ((Coord, Maybe Cell, Int) -> Coord)
-> ((Coord, Maybe Cell, Int) -> [(Coord, Maybe Cell, Int)])
-> (Coord, Maybe Cell, Int)
-> [(Coord, Maybe Cell, Int)]
forall r a. Ord r => (a -> r) -> (a -> [a]) -> a -> [a]
bfsOn (\(Coord
p,Maybe Cell
_,Int
_)->Coord
p) (Coord, Maybe Cell, Int) -> [(Coord, Maybe Cell, Int)]
forall {c} {a}.
Num c =>
(Coord, Maybe a, c) -> [(Coord, Maybe Cell, c)]
step (Coord
start, Cell -> Maybe Cell
forall a. a -> Maybe a
Just Cell
startCell, Int
0)
]
where
step :: (Coord, Maybe a, c) -> [(Coord, Maybe Cell, c)]
step (Coord
here, Maybe a
hereCell, c
n)
| Coord
here Coord -> Coord -> Bool
forall a. Eq a => a -> a -> Bool
/= Coord
start Bool -> Bool -> Bool
&& Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
hereCell = []
| Bool
otherwise =
[ (Coord
there, Maybe Cell
thereCell, c
nc -> c -> c
forall a. Num a => a -> a -> a
+c
1)
| Coord
there <- Coord -> [Coord]
cardinal Coord
here
, let char :: Char
char = UArray Coord Char
world UArray Coord Char -> Coord -> Char
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Coord
there
, let thereCell :: Maybe Cell
thereCell = Char -> Maybe Cell
charToCell Char
char
, Char
char Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'#'
]
data AllKeys = AllKeys
{ AllKeys -> IntSet
akKeys :: !IntSet
, AllKeys -> Set Coord
akLocations :: !(Set Coord)
}
deriving (Eq AllKeys
Eq AllKeys =>
(AllKeys -> AllKeys -> Ordering)
-> (AllKeys -> AllKeys -> Bool)
-> (AllKeys -> AllKeys -> Bool)
-> (AllKeys -> AllKeys -> Bool)
-> (AllKeys -> AllKeys -> Bool)
-> (AllKeys -> AllKeys -> AllKeys)
-> (AllKeys -> AllKeys -> AllKeys)
-> Ord AllKeys
AllKeys -> AllKeys -> Bool
AllKeys -> AllKeys -> Ordering
AllKeys -> AllKeys -> AllKeys
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: AllKeys -> AllKeys -> Ordering
compare :: AllKeys -> AllKeys -> Ordering
$c< :: AllKeys -> AllKeys -> Bool
< :: AllKeys -> AllKeys -> Bool
$c<= :: AllKeys -> AllKeys -> Bool
<= :: AllKeys -> AllKeys -> Bool
$c> :: AllKeys -> AllKeys -> Bool
> :: AllKeys -> AllKeys -> Bool
$c>= :: AllKeys -> AllKeys -> Bool
>= :: AllKeys -> AllKeys -> Bool
$cmax :: AllKeys -> AllKeys -> AllKeys
max :: AllKeys -> AllKeys -> AllKeys
$cmin :: AllKeys -> AllKeys -> AllKeys
min :: AllKeys -> AllKeys -> AllKeys
Ord, AllKeys -> AllKeys -> Bool
(AllKeys -> AllKeys -> Bool)
-> (AllKeys -> AllKeys -> Bool) -> Eq AllKeys
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: AllKeys -> AllKeys -> Bool
== :: AllKeys -> AllKeys -> Bool
$c/= :: AllKeys -> AllKeys -> Bool
/= :: AllKeys -> AllKeys -> Bool
Eq, Int -> AllKeys -> ShowS
[AllKeys] -> ShowS
AllKeys -> String
(Int -> AllKeys -> ShowS)
-> (AllKeys -> String) -> ([AllKeys] -> ShowS) -> Show AllKeys
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> AllKeys -> ShowS
showsPrec :: Int -> AllKeys -> ShowS
$cshow :: AllKeys -> String
show :: AllKeys -> String
$cshowList :: [AllKeys] -> ShowS
showList :: [AllKeys] -> ShowS
Show)
allKeys ::
UArray Coord Char ->
[Coord] ->
Int
allKeys :: UArray Coord Char -> [Coord] -> Int
allKeys UArray Coord Char
world [Coord]
start =
[(AllKeys, Int)] -> Int
forall {c}. [(AllKeys, c)] -> c
select ([(AllKeys, Int)] -> Int) -> [(AllKeys, Int)] -> Int
forall a b. (a -> b) -> a -> b
$ (AllKeys -> [AStep AllKeys]) -> AllKeys -> [(AllKeys, Int)]
forall a. Ord a => (a -> [AStep a]) -> a -> [(a, Int)]
astar AllKeys -> [AStep AllKeys]
stepAK (AllKeys -> [(AllKeys, Int)]) -> AllKeys -> [(AllKeys, Int)]
forall a b. (a -> b) -> a -> b
$ IntSet -> Set Coord -> AllKeys
AllKeys IntSet
IntSet.empty (Set Coord -> AllKeys) -> Set Coord -> AllKeys
forall a b. (a -> b) -> a -> b
$ [Coord] -> Set Coord
forall a. Ord a => [a] -> Set a
Set.fromList [Coord]
start
where
keyN :: Int
keyN = (Char -> Bool) -> String -> Int
forall (f :: * -> *) a. Foldable f => (a -> Bool) -> f a -> Int
countBy Char -> Bool
isLower (UArray Coord Char -> String
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
elems UArray Coord Char
world)
done :: (AllKeys, b) -> Bool
done (AllKeys, b)
s = IntSet -> Int
IntSet.size (AllKeys -> IntSet
akKeys ((AllKeys, b) -> AllKeys
forall a b. (a, b) -> a
fst (AllKeys, b)
s)) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
keyN
select :: [(AllKeys, c)] -> c
select = (AllKeys, c) -> c
forall a b. (a, b) -> b
snd ((AllKeys, c) -> c)
-> ([(AllKeys, c)] -> (AllKeys, c)) -> [(AllKeys, c)] -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(AllKeys, c)] -> (AllKeys, c)
forall a. HasCallStack => [a] -> a
head ([(AllKeys, c)] -> (AllKeys, c))
-> ([(AllKeys, c)] -> [(AllKeys, c)])
-> [(AllKeys, c)]
-> (AllKeys, c)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((AllKeys, c) -> Bool) -> [(AllKeys, c)] -> [(AllKeys, c)]
forall a. (a -> Bool) -> [a] -> [a]
filter (AllKeys, c) -> Bool
forall {b}. (AllKeys, b) -> Bool
done
paths :: Map Coord [(Coord, Cell, Int)]
paths = UArray Coord Char -> Map Coord [(Coord, Cell, Int)]
extractGraph UArray Coord Char
world
stepAK :: AllKeys -> [AStep AllKeys]
stepAK AllKeys{Set Coord
IntSet
akKeys :: AllKeys -> IntSet
akLocations :: AllKeys -> Set Coord
akKeys :: IntSet
akLocations :: Set Coord
..} =
[ AStep {
astepNext :: AllKeys
astepNext = IntSet -> Set Coord -> AllKeys
AllKeys (Int -> IntSet -> IntSet
IntSet.insert Int
k IntSet
akKeys)
(Coord -> Set Coord -> Set Coord
forall a. Ord a => a -> Set a -> Set a
Set.insert Coord
loc (Coord -> Set Coord -> Set Coord
forall a. Ord a => a -> Set a -> Set a
Set.delete Coord
who Set Coord
akLocations)),
astepCost :: Int
astepCost = Int
cost,
astepHeuristic :: Int
astepHeuristic = Int
0 }
| Coord
who <- Set Coord -> [Coord]
forall a. Set a -> [a]
Set.toList Set Coord
akLocations
, let Just Cell
whoCell = Char -> Maybe Cell
charToCell (UArray Coord Char
world UArray Coord Char -> Coord -> Char
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Coord
who)
, (Coord
loc, Int
k, Int
cost) <- Map Coord [(Coord, Cell, Int)]
-> Coord -> Cell -> IntSet -> [(Coord, Int, Int)]
nextKey Map Coord [(Coord, Cell, Int)]
paths Coord
who Cell
whoCell IntSet
akKeys
]
nextKey ::
Map Coord [(Coord, Cell, Int)] ->
Coord ->
Cell ->
IntSet ->
[(Coord, Int, Int)]
nextKey :: Map Coord [(Coord, Cell, Int)]
-> Coord -> Cell -> IntSet -> [(Coord, Int, Int)]
nextKey Map Coord [(Coord, Cell, Int)]
paths Coord
start Cell
startCell IntSet
keys =
[ (Coord
here, Int
k, Int
cost)
| ((Coord
here, Key Int
k), Int
cost) <- ((Coord, Cell) -> Coord)
-> ((Coord, Cell) -> [AStep (Coord, Cell)])
-> (Coord, Cell)
-> [((Coord, Cell), Int)]
forall b a.
Ord b =>
(a -> b) -> (a -> [AStep a]) -> a -> [(a, Int)]
astarOn (Coord, Cell) -> Coord
forall a b. (a, b) -> a
fst (Coord, Cell) -> [AStep (Coord, Cell)]
step (Coord
start,Cell
startCell) ]
where
step :: (Coord, Cell) -> [AStep (Coord, Cell)]
step (Coord
here, Cell
hereCell) =
[ (Coord, Cell) -> Int -> Int -> AStep (Coord, Cell)
forall a. a -> Int -> Int -> AStep a
AStep (Coord
loc, Cell
cell) Int
cost Int
0
| case Cell
hereCell of
Key Int
k -> Int -> IntSet -> Bool
IntSet.member Int
k IntSet
keys
Cell
_ -> Bool
True
, (Coord
loc, Cell
cell, Int
cost) <- Map Coord [(Coord, Cell, Int)]
paths Map Coord [(Coord, Cell, Int)] -> Coord -> [(Coord, Cell, Int)]
forall k a. Ord k => Map k a -> k -> a
Map.! Coord
here
, case Cell
cell of
Gate Int
i -> Int -> IntSet -> Bool
IntSet.member Int
i IntSet
keys
Cell
_ -> Bool
True
]
keysSSP ::
UArray Coord Char ->
Map Coord [(Coord, x, Int)] ->
Map (Coord, Coord) Int
keysSSP :: forall x.
UArray Coord Char
-> Map Coord [(Coord, x, Int)] -> Map (Coord, Coord) Int
keysSSP UArray Coord Char
world Map Coord [(Coord, x, Int)]
direct = ((Coord, Coord) -> Int -> Bool)
-> Map (Coord, Coord) Int -> Map (Coord, Coord) Int
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (Coord, Coord) -> Int -> Bool
forall {p}. (Coord, Coord) -> p -> Bool
scrub ((Map (Coord, Coord) Int -> Coord -> Map (Coord, Coord) Int)
-> Map (Coord, Coord) Int -> [Coord] -> Map (Coord, Coord) Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map (Coord, Coord) Int -> Coord -> Map (Coord, Coord) Int
forall {a}.
(Num a, Ord a) =>
Map (Coord, Coord) a -> Coord -> Map (Coord, Coord) a
addGen Map (Coord, Coord) Int
gen0 [Coord]
ks)
where
scrub :: (Coord, Coord) -> p -> Bool
scrub (Coord
c1,Coord
c2) p
_
| Just Key{} <- Char -> Maybe Cell
charToCell (UArray Coord Char
world UArray Coord Char -> Coord -> Char
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Coord
c1)
, Just Key{} <- Char -> Maybe Cell
charToCell (UArray Coord Char
world UArray Coord Char -> Coord -> Char
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Coord
c2) = Bool
True
| Bool
otherwise = Bool
False
ks :: [Coord]
ks = Map Coord [(Coord, x, Int)] -> [Coord]
forall k a. Map k a -> [k]
Map.keys Map Coord [(Coord, x, Int)]
direct
gen0 :: Map (Coord, Coord) Int
gen0 = [((Coord, Coord), Int)] -> Map (Coord, Coord) Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [ ((Coord
src,Coord
dst), Int
cost)
| (Coord
src,[(Coord, x, Int)]
dsts) <- Map Coord [(Coord, x, Int)] -> [(Coord, [(Coord, x, Int)])]
forall k a. Map k a -> [(k, a)]
Map.toList Map Coord [(Coord, x, Int)]
direct
, (Coord
dst,x
_,Int
cost) <- [(Coord, x, Int)]
dsts
]
mkCost :: Maybe a -> Maybe a -> Maybe a -> [a]
mkCost Maybe a
Nothing (Just a
ik) (Just a
kj) = [a
ika -> a -> a
forall a. Num a => a -> a -> a
+a
kj]
mkCost (Just a
ij) (Just a
ik) (Just a
kj) = [a -> a -> a
forall a. Ord a => a -> a -> a
min a
ij (a
ika -> a -> a
forall a. Num a => a -> a -> a
+a
kj)]
mkCost (Just a
ij) Maybe a
_ Maybe a
_ = [a
ij]
mkCost Maybe a
Nothing Maybe a
_ Maybe a
_ = []
addGen :: Map (Coord, Coord) a -> Coord -> Map (Coord, Coord) a
addGen Map (Coord, Coord) a
acc Coord
k = [((Coord, Coord), a)] -> Map (Coord, Coord) a
forall k a. Ord k => [(k, a)] -> Map k a
MapStrict.fromList
[ ((Coord
i,Coord
j), a
cost)
| Coord
i <- [Coord]
ks
, Coord
j <- [Coord]
ks
, a
cost <- Maybe a -> Maybe a -> Maybe a -> [a]
forall {a}. (Num a, Ord a) => Maybe a -> Maybe a -> Maybe a -> [a]
mkCost ((Coord, Coord) -> Map (Coord, Coord) a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Coord
i,Coord
j) Map (Coord, Coord) a
acc)
((Coord, Coord) -> Map (Coord, Coord) a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Coord
i,Coord
k) Map (Coord, Coord) a
acc)
((Coord, Coord) -> Map (Coord, Coord) a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Coord
k,Coord
j) Map (Coord, Coord) a
acc)
]