day 17 writeup

src/Javran/AdventOfCode/Y2023/Day17.hs

@@ -15,42 +15,74 @@ data Day17 deriving (Generic)
 
 type Dims = (Int, Int) -- rows then cols
 
-parseGraph :: [String] -> (Dims, Coord -> Maybe Int)
-parseGraph raw = ((rows, cols), getV)
-  where
-    rows = length raw
-    cols = length (head raw)
-    bounds = ((0, 0), (rows - 1, cols - 1))
+{-
+  [General notes]
 
-    getV coord =
-      vs VU.! index bounds coord <$ guard (inRange bounds coord)
+  Nothing fancy, just your usual A* search.
 
-    vs :: VU.Vector Int
-    vs = VU.fromList $ fmap (\v -> ord v - ord '0') $ concat raw
+  It seems to be the case that input digits do not contain `0`, this means
+  every step is attached with a cost of at least 1 - this is a nice property
+  as we can just let heuristic function be the manhattan distance from current position
+  to target position. In addition, as this heuristic function is underestimating the cost,
+  we always reach the optimal solution first, if it can be found.
 
-{-
-  TODO:
+  For part 1, we encode constraint in a state (which is the `CSt` type below),
+  So that we keep track of how many consecutive moves we are allowed to make
+  by representing this number as a "move budget", when it reaches zero,
+  one can no longer go in the same direction.
 
-  - I'm wondering if this is yet another boring a*-search-with-a-twist-thingy:
+  For part 2, we can't do the same thing as easy as part 1 - this is because
+  ultra crusible does not have the flexibility to "stop after one step",
+  as it goes 4~10 steps all the way, meaning if we were to encode the state same way:
 
-    for the 2d plane, a state can be encoded:
+  + there might be illegal states (say only moves twice in same direction)
+    that we need ways to get rid of
 
-    (coord, dir moving into it, budget for going straight)
+  + we might have too many states in the search to keep track of (because now
+    we have a budget of 10 consecutive moves to encode)
 
-    and each of those states will have a clear "best value" as we conduct the search.
+  However we can do things slightly differently to resolve this issue,
+  which is to encode this constraint by state transition
+  (this is the `nextsOf2` function below):
 
-  - A few things find while looking at input data:
+  We no longer have a "move budget", but instead:
 
-    + no `0` in the input, so we can always use manhattan distance to bottom right
-      as an overestimation.
+  + state transition "jumps" a state 4~10 steps in all directions directly
 
-    + if we text search `1`, `2`, ... and so on in actual input we can see
-      a pattern that lower values are scattered outside and higher values are around
-      the center. i'm not sure if this observation would lead us anywhere,
-      but i feel like this worth taking a note for some reason.
+  + previous direction is still recorded. but state transition no longer
+    attempts to move in the same (or opposite) direction of previous one.
+
+  By doing so we also make sure when we reach final coordinate the ultra crusible
+  is always ready to be stopped.
+
+  [An interesting but probably irrelevant observation]
+
+  If we text search `1`, `2`, ... and so on in actual input we can see
+  a pattern that lower values are scattered outside and higher values are around
+  the center. This gives us a slight hint about the strategy behind how those
+  inputs are generated: having high values in the center will encourage path
+  of low cost to "go around" so we don't get a straight permutation of moving
+  either right or down all the way to target location.
 
  -}
 
+{-
+  Parses input graph and gets back its dimensions and individual values
+  (in the form of a safe getter).
+ -}
+parseGraph :: [String] -> (Dims, Coord -> Maybe Int)
+parseGraph raw = ((rows, cols), getV)
+  where
+    rows = length raw
+    cols = length (head raw)
+    bounds = ((0, 0), (rows - 1, cols - 1))
+
+    getV coord =
+      vs VU.! index bounds coord <$ guard (inRange bounds coord)
+
+    vs :: VU.Vector Int
+    vs = VU.fromList $ fmap (\v -> ord v - ord '0') $ concat raw
+
 {-
   Crusible state for the search:
 
@@ -73,79 +105,136 @@ parseGraph raw = ((rows, cols), getV)
  -}
 type CSt = (Coord, Dir, Int)
 
+{-
+  Part 1 state transitiion. Returns list of next states together with
+  its corresponding cost.
+ -}
 nextsOf :: (Coord -> Maybe Int) -> CSt -> [(CSt, Int)]
 nextsOf getCost (cur, prevDir, budget) = do
   d <- allDirs
   let next = applyDir d cur
   Just cost <- [getCost next]
   if d == prevDir
     then do
+      -- it costs budget to go the same direction as previous one
       guard $ budget > 0
       pure ((next, d, budget - 1), cost)
     else do
+      -- turning directly back is not allowed.
       guard $ d /= oppositeDir prevDir
+      -- it resets budget if we choose to go a different direction.
       pure ((next, d, 2), cost)
 
 {-
-  TODO: p2 notes
-
-  We probably need to approach p2 in a slightly different way, because
-  ultra crusible does not have the flexibility to "stop after one step",
-  as it goes 4~10 steps all the way.
-
-  it is probably easier to take care of this constraint *not* in the state
-  but in state transition - we no longer have a "move budget" but instead
-  state transition "jumps" a state 4~10 steps directly - by doing so
-  we also make sure when we reach final coordinate the ultra crusible
-  is always ready to be stopped.
+  Encodes how an ultra crusible moves without changing direction.
+  Returns coordinates that it is stopped on and the total cost of moving to there.
 
+  This serves as a building block for state transition for part 2.
  -}
-
 ultraCrusibleMove :: (Coord -> Maybe Int) -> Dir -> Coord -> [(Coord, Int)]
 ultraCrusibleMove getCost d p0 = do
+  {-
+    Admittedly there are some amount of repetition and imperative style
+    to this implementation, but I'd argue:
+
+    - the action of taking certain steps is imperative in nature
+
+    - it's important to lay down these steps explicitly than trying to abstract over it,
+      as it simply doesn't help with understanding this imperative process.
+
+   -}
+  let
+    step :: Coord -> [(Coord, Int)]
+    step a = do
+      let a' = applyDir d a
+      Just c <- [getCost a']
+      pure (a', c)
+
+  {-
+    This part takes care of moving 4 steps in a direction,
+    which is mandatory before the ultra crusible can be stopped.
+   -}
   (p1, acc1) <-
     fix
       ( \loop n pNow acc ->
-          if n > (0 :: Int)
+          if n > 0
             then do
               (p', extra) <- step pNow
-              loop (n - 1) p' (acc + extra)
+              loop (n - 1 :: Int) p' (acc + extra)
             else pure (pNow, acc)
       )
       4
       p0
       0
+  {- Now an extra of 0~6 steps can be taken. -}
   fix
     ( \loop n pNow acc ->
+        -- stop at any time, otherwise take an extra step.
         (pNow, acc)
           : ( do
-                guard $ n > (0 :: Int)
+                guard $ n > 0
                 (p', extra) <- step pNow
-                loop (n - 1) p' (acc + extra)
+                loop (n - 1 :: Int) p' (acc + extra)
             )
     )
     6
     p1
     acc1
-  where
-    step :: Coord -> [(Coord, Int)]
-    step a = do
-      let a' = applyDir d a
-      Just c <- [getCost a']
-      pure (a', c)
 
 type CSt2 = (Coord, Dir) -- now just coord and direction of last move
 
-nextsOf2 :: (Coord -> Maybe Int) -> Maybe Dir -> Coord -> [] (CSt2, Int)
+{-
+  Part 2 state transition differs from part 1 in that after take a (non-determinstic) move
+  in one direction, we are no longer allowed to go in that previous direction
+  (or opposite of previous direction).
+ -}
+nextsOf2 :: (Coord -> Maybe Int) -> Maybe Dir -> Coord -> [(CSt2, Int)]
 nextsOf2 getCost mPrevDir cur = do
-  d <- allDirs
-  guard $ maybe True (\d' -> d' /= d && d' /= oppositeDir d) mPrevDir
+  d <- case mPrevDir of
+    Nothing -> allDirs
+    Just prevDir ->
+      -- not allowed to go straight or backwards, leaving only left and right turns.
+      [turnLeft prevDir, turnRight prevDir]
   (next, cost) <- ultraCrusibleMove getCost d cur
   pure ((next, d), cost)
 
+{-
+  See: https://en.wikipedia.org/wiki/A*_search_algorithm
+
+  `aStar estToFinal getNexts dists q` searches for the least cost (or heat loss in this problem).
+
+  - estToFinal: an estimated cost from current state to final state.
+
+    In this problem we simply use manhattan distance from current coordinate to the final one,
+    as this is an underestimation, we always reach optimal solution when we can find a solution.
+
+  - getNexts: state transition function.
+
+  - dists: dists[n] is the cost of the cheapest path from start to n currently known.
+    In Wikipedia article this is referred to as "gScore"
+
+  - q: of type `PSQ st (Arg fScore dist)` the priority queue.
+
+    + `st` is the current state, same as keys of `dists`: when we are computing next states given current state,
+      we "throw away" states that we both (1) have seen before (2) this new expansion does not improve current known cost.
+      (note that we are not necessarily using the "search" feature of PSQ in a direct way)
+
+    + `fScore` is the current best guess of total cost going from start, via `st`, to final goal state.
+
+    + `dist` is the cost already incurred going from start state to `st`.
+
+      If we can ensure that `estToFinal` never returns 0 on a non-goal state (it is the case for this problem),
+      the goal state is found when `dist == fScore`.
+
+    Here note that values are `Arg fScore dist`, meaning only `fScore` is being used to sort out priority,
+    and `dist` is just an extra piece of information that we attach to current state.
+    I find it, at least in theory, slightly more efficient than having to keep `fScore` as a dictionary that
+    has to be updated, passed around, and looked up.
+
+ -}
 aStar :: forall st. Ord st => (st -> Int) -> (st -> [(st, Int)]) -> M.Map st Int -> PQ.PSQ st (Arg Int Int) -> Int
 aStar estToFinal getNexts =
-  fix \go (dists :: M.Map st Int) q0 -> case PQ.minView q0 of
+  fix \go dists q0 -> case PQ.minView q0 of
     Nothing -> error "queue exhausted"
     Just (cur PQ.:-> Arg fScore dist, q1) ->
       if fScore == dist
@@ -155,6 +244,7 @@ aStar estToFinal getNexts =
                 (next, moveCost) <- getNexts cur
                 let dist' = dist + moveCost
                     fScore' = dist' + estToFinal next
+                -- reject next states that are not improving current known costs.
                 guard $ maybe True (dist' <) (dists M.!? next)
                 pure (next, dist', Arg fScore' dist')
               q2 = foldr (\(next, _, p) -> PQ.insert next p) q1 nexts
@@ -164,16 +254,49 @@ aStar estToFinal getNexts =
 instance Solution Day17 where
   solutionRun _ SolutionContext {getInputS, answerShow} = do
     (dims, getV) <- parseGraph . lines <$> getInputS
-    let (rows, cols) = dims
-        estToFinal :: Coord -> Int
-        estToFinal = manhattan (rows - 1, cols - 1)
-
-    answerShow $ aStar (estToFinal . (\(p, _, _) -> p)) (nextsOf getV) M.empty (PQ.singleton ((0, 0), D, 3) (Arg (estToFinal (0, 0)) 0))
-
-    let initNexts :: [] (CSt2, Int)
-        initNexts = nextsOf2 getV Nothing (0, 0)
+    let
+      (rows, cols) = dims
+      start = (0, 0)
+      goal = (rows - 1, cols - 1)
+      estToFinal :: Coord -> Int
+      estToFinal = manhattan goal
+
+    {-
+      Here we enqueue an illegal state as starting point
+      because a state must have a previous state.
+
+      By using `3` as budget, it is allowed to take an extra to compensate
+      for the fake "previous direction".
+
+      We could put it in initial `dists`, I don't feel it's necessary -
+      it would not be looked up during the search anyway.
+
+     -}
+    answerShow $
+      aStar @CSt
+        (estToFinal . (\(p, _, _) -> p))
+        (nextsOf getV)
+        M.empty
+        (PQ.singleton (start, D, 3) (Arg (estToFinal start) 0))
+
+    do
+      {-
+        Here we face a similar problem as part 1:
+        to construct `CSt2` we need to give it a previous direction.
+
+        This time we solve it by taking the initial step outside of search algorithm,
+        and "feed those back" to kick start it.
+       -}
+      let
+        initNexts :: [(CSt2, Int)]
+        initNexts = nextsOf2 getV Nothing start
         initQ = PQ.fromList do
           (s@(p, _), dist) <- initNexts
           pure (s PQ.:-> Arg (dist + estToFinal p) dist)
-        getNexts2 (p, d) = nextsOf2 getV (Just d) p
-    answerShow $ aStar (estToFinal . fst) getNexts2 (M.fromList initNexts) initQ
+
+      answerShow $
+        aStar @CSt2
+          (estToFinal . fst)
+          (\(p, d) -> nextsOf2 getV (Just d) p)
+          (M.fromList initNexts)
+          initQ