Day 18 part 2

1 files changed, 67 insertions, 34 deletions

diff --git a/2019/18.hs b/2019/18.hs
index a60421a..c84f710 100644
--- a/2019/18.hs
+++ b/2019/18.hs
@@ -21,6 +21,11 @@ import qualified SmallIntSet as SIS
 import SmallIntSet (SmallIntSet)
 
 
+replaceAtIndex :: Int -> a -> [a] -> [a]
+replaceAtIndex i _ [] = error ("Index is " ++ show i ++ " items past end in replaceAtIndex")
+replaceAtIndex 0 val (_:xs) = val : xs
+replaceAtIndex i val (x:xs) = x : replaceAtIndex (i-1) val xs
+
 -- Considers a distance of '-1' to mean 'unconnected'.
 -- Applies Floyd-Warshall.
 transitiveClosure :: [Int] -> A.UArray (Int, Int) Int -> A.UArray (Int, Int) Int
@@ -57,30 +62,45 @@ reachableFrom bd startPos = go 0 (Set.singleton startPos) (Set.singleton startPo
                then result'
                else go (dist + 1) (seen <> Set.fromList boundary') (Set.fromList frees) result'
 
--- [0  1  26   27 52]
--- [@, a...z,  A...Z]
-data Implicit = Implicit (A.Array Int [Int])        -- edge list
+-- [0  25   26 51  52  ]
+-- [a...z,  A...Z, @...]
+data Implicit = Implicit Int                        -- number of starting positions
+                         (A.Array Int [Int])        -- edge list
                          (A.UArray (Int, Int) Int)  -- distance matrix, with closure taken
   deriving (Show)
 
-implicitGraph :: Map Pos Char -> Pos -> Implicit
-implicitGraph bd startPos =
-    let posGraph = fst (go startPos Map.empty Set.empty)
+codeIsStart, codeIsLower, codeIsUpper :: Int -> Bool
+codeIsStart n = n >= 2 * 26
+codeIsLower n = n < 26
+codeIsUpper n = 26 <= n && n < 2 * 26
+
+codeToLower :: Int -> Int
+codeToLower n = n - 26     -- assumes codeIsUpper
+
+implicitGraph :: Map Pos Char -> [Pos] -> Implicit
+implicitGraph bd startPositions =
+    let nStart = length startPositions
+        nNodes = 2 * 26 + nStart
+        posGraph = fst (goMultiple startPositions Map.empty Set.empty)
         mapGraph = Map.mapKeys (bd Map.!) (Map.map (Map.mapKeys (bd Map.!)) posGraph)
-        charToNode '@' = 0
-        charToNode c | isLower c = 1 + ord c - ord 'a'
-                     | isUpper c = 1 + 26 + ord c - ord 'A'
+        charToNode '@' = charToNode '1'
+        charToNode c | isLower c = ord c - ord 'a'
+                     | isUpper c = 26 + ord c - ord 'A'
+                     | isDigit c = 26 + 26 + ord c - ord '1'
                      | otherwise = undefined
-        arrGraph = A.accumArray (const id) [] (0, 2 * 26)
+        arrGraph = A.accumArray (const id) [] (0, nNodes - 1)
                         [(charToNode from, map charToNode (Map.keys tomap))
                         | (from, tomap) <- Map.assocs mapGraph]
-        distArr = A.accumArray (const id) (-1) ((0, 0), (2 * 26, 2 * 26))
+        distArr = A.accumArray (const id) (-1) ((0, 0), (nNodes - 1, nNodes - 1))
                         [((charToNode from, charToNode to), dist)
                         | (from, tomap) <- Map.assocs mapGraph
                         , (to, dist) <- Map.assocs tomap]
         nodeList = map charToNode (Map.keys mapGraph)
-    in Implicit arrGraph (transitiveClosure nodeList distArr)
+    in Implicit nStart arrGraph (transitiveClosure nodeList distArr)
   where
+    goMultiple :: [Pos] -> Map Pos (Map Pos Int) -> Set Pos -> (Map Pos (Map Pos Int), Set Pos)
+    goMultiple curPoses graph seen = foldl' (\(gr, sn) node -> go node gr sn) (graph, seen) curPoses
+
     go :: Pos -> Map Pos (Map Pos Int) -> Set Pos -> (Map Pos (Map Pos Int), Set Pos)
     go curPos graph seen
         | curPos `Set.member` seen = (graph, seen)
@@ -89,15 +109,15 @@ implicitGraph bd startPos =
                 newNodes = Map.keysSet reach Set.\\ seen
                 graph' = Map.insert curPos reach graph
                 seen' = Set.insert curPos seen
-            in Set.foldl' (\(gr, sn) node -> go node gr sn) (graph', seen') newNodes
+            in goMultiple (Set.toList newNodes) graph' seen'
 
 reachable :: Implicit -> SmallIntSet -> Int -> IntMap Int
-reachable (Implicit graph distarr) keys start = snd (go 0 (SIS.singleton start) start IntMap.empty)
+reachable (Implicit _ graph distarr) keys start = snd (go 0 (SIS.singleton start) start IntMap.empty)
   where
     go dist seen at result =
-        let nexts = filter (\c -> c `SIS.notMember` seen && (c <= 26 || (c - 26) `SIS.member` keys))
+        let nexts = filter (\c -> c `SIS.notMember` seen && (codeIsStart c || codeIsLower c || codeToLower c `SIS.member` keys))
                            (graph A.! at)
-            (nextPearls, nextNonpearls) = partition (\c -> 0 < c && c <= 26 && c `SIS.notMember` keys) nexts
+            (nextPearls, nextNonpearls) = partition (\c -> codeIsLower c && c `SIS.notMember` keys) nexts
             result' = result <> IntMap.fromList [(c, dist + distarr A.! (at, c)) | c <- nextPearls]
             seen' = seen <> SIS.fromList nexts
         in -- trace ("reachable-go at=" ++ show at ++ " dist=" ++ show dist ++ " nexts=" ++ show nexts ++ " (allnexts " ++ show (graph A.! at) ++ ")") $
@@ -106,37 +126,41 @@ reachable (Implicit graph distarr) keys start = snd (go 0 (SIS.singleton start)
                else foldl (\(sn, rs) c -> go (dist + distarr A.! (at, c)) sn c rs) (seen', result') nextNonpearls
 
 searchBFS :: SmallIntSet -> Implicit -> (Int, [Int])
-searchBFS allKeys implicit@(Implicit _ distarr) = go 0 (Set.singleton (heuristic 0 SIS.empty, 0, 0, SIS.empty, [])) Map.empty
+searchBFS allKeys implicit@(Implicit nstart _ distarr) =
+    let startNodes = [52 + i | i <- [0..nstart-1]]
+    in go 0 (Set.singleton (heuristic startNodes SIS.empty, 0, startNodes, SIS.empty, [])) Map.empty
   where
-    -- pqueue: f-val, distance, node, keys, key order
-    -- visited: node, keys => distance
-    go :: Int -> Set (Int, Int, Int, SmallIntSet, [Int]) -> Map (Int, SmallIntSet) Int -> (Int, [Int])
+    -- pqueue: f-val, distance, nodes, keys, key order
+    -- visited: nodes, keys => distance
+    go :: Int -> Set (Int, Int, [Int], SmallIntSet, [Int]) -> Map ([Int], SmallIntSet) Int -> (Int, [Int])
     go ctr pqueue visited =
-        let ((heurval, dist, curnode, keys, keyorder), newpqueue) = Set.deleteFindMin pqueue
-            reach = reachable implicit keys curnode
-            nextStates = [(dist + stepDist + heuristic stepC stepKeys, dist + stepDist, stepC, stepKeys, stepC : keyorder)
-                         | (stepC, stepDist) <- IntMap.assocs reach
+        let ((heurval, dist, curnodes, keys, keyorder), newpqueue) = Set.deleteFindMin pqueue
+            reachLists = [reachable implicit keys node | node <- curnodes]
+            nextStates = [(dist + stepDist + heuristic stepNodes stepKeys, dist + stepDist, stepNodes, stepKeys, stepC : keyorder)
+                         | (robotIdx, reach) <- zip [0..] reachLists
+                         , (stepC, stepDist) <- IntMap.assocs reach
                          , let stepKeys = SIS.insert stepC keys
+                               stepNodes = replaceAtIndex robotIdx stepC curnodes
                            -- check that this next state is actually better than we've seen before
-                         , maybe True (dist + stepDist <) (Map.lookup (stepC, stepKeys) visited)]
-            visited' = Map.insert (curnode, keys) dist visited
+                         , maybe True (dist + stepDist <) (Map.lookup (stepNodes, stepKeys) visited)]
+            visited' = Map.insert (curnodes, keys) dist visited
             pqueue' = newpqueue <> Set.fromList nextStates
             result = 
-               if IntMap.null reach
+               if all IntMap.null reachLists
                    then if heurval == dist then (dist, keyorder) else error ("heurval - dist = " ++ show (heurval - dist) ++ " in terminal state!")
                    else go (ctr + 1) pqueue' visited'
-        in -- (if ctr `rem` 20000 == 0 || IntMap.null reach
+        in -- (if ctr `rem` 20000 == 0 || all IntMap.null reachLists
            --      then trace ("go #pqueue=" ++ show (Set.size pqueue) ++ " #visited=" ++ show (Map.size visited)
-           --                  ++ "  curnode=" ++ show curnode ++ " dist=" ++ show dist ++ " heurval=" ++ show heurval ++ " keys=" ++ show keyorder
+           --                  ++ "  curnodes=" ++ show curnodes ++ " dist=" ++ show dist ++ " heurval=" ++ show heurval ++ " keys=" ++ show keyorder
            --                  {- ++ "  next->" ++ show nextStates -})
            --      else id)
            result
 
-    heuristic :: Int -> SmallIntSet -> Int
-    heuristic _curnode keys =
+    heuristic :: [Int] -> SmallIntSet -> Int
+    heuristic _curnodes keys =
         let remainKeys = allKeys SIS.\\ keys
             allDists = [distarr A.! (x, y) | x:xs <- tails (SIS.toList remainKeys), y <- xs]
-            distLowerBound = sum (take (SIS.size remainKeys - 1) (sort allDists))
+            distLowerBound = sum (take (SIS.size remainKeys - 1) (sort (filter (/= -1) allDists)))
         in distLowerBound
 
 main :: IO ()
@@ -145,6 +169,15 @@ main = do
     let bd = Map.fromList [((x, y), c) | (y, row) <- zip [0..] stringbd, (x, c) <- zip [0..] row]
         startpos = fromJust (lookup '@' (map (\(x,y) -> (y,x)) (Map.assocs bd)))
 
-    let imgraph = implicitGraph bd startpos
-        allKeys = SIS.fromList [ord c - ord 'a' + 1 | c <- Map.elems bd, isLower c]
+    let imgraph = implicitGraph bd [startpos]
+        allKeys = SIS.fromList [ord c - ord 'a' | c <- Map.elems bd, isLower c]
     print (fst (searchBFS allKeys imgraph))
+
+    let (sx, sy) = startpos
+        bd2 = Map.unionWith (const id) bd
+                            (Map.fromList [((sx-1,sy-1), '1'), ((sx,sy-1), '#'), ((sx+1,sy-1), '2')
+                                          ,((sx-1,sy  ), '#'), ((sx,sy  ), '#'), ((sx+1,sy  ), '#')
+                                          ,((sx-1,sy+1), '3'), ((sx,sy+1), '#'), ((sx+1,sy+1), '4')])
+        imgraph2 = implicitGraph bd2 [(sx-1,sy-1), (sx+1,sy-1), (sx-1,sy+1), (sx+1,sy+1)]
+
+    print (fst (searchBFS allKeys imgraph2))