day15: part2 is 2545x faster

Author: siraben

day15.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE TupleSections #-}
 
+import Criterion.Main
 import Data.Function
 import qualified Data.IntMap as IM
 import Data.IntMap.Strict (IntMap)
@@ -12,21 +13,26 @@ type Grid = IntMap Int
 
 type Pt = Int
 
+factor :: Int
+factor = 10000
+
 conv :: (Int, Int) -> Pt
-conv (x, y) = 10000 * x + y
+conv (x, y) = factor * x + y
 
 from :: Pt -> (Int, Int)
-from n = (n `div` 10000, n `mod` 10000)
+from n = (n `div` factor, n `mod` factor)
 
 toGrid :: [[Int]] -> Grid
 toGrid = IM.fromList . concat . zipWith (\x l -> zipWith (curry (\((x, a), y) -> (conv (x, y), a))) (map (x,) l) [1 ..]) [1 ..]
 
+toGrid2 :: [[Int]] -> Grid
 toGrid2 = toGrid . embig
 
 succW :: Int -> Int
 succW n = max ((n + 1) `mod` 10) 1
 
-embig l = concat $ take 5 $ iterate (map (map succW)) [concat $ take 5 (iterate (map succW) r) | r <- l]
+embig :: [[Int]] -> [[Int]]
+embig = concat . take 5 . iterate (map (map succW)) . map (concat . take 5 . iterate (map succW))
 
 {-
  1  function Dijkstra(Graph, source):
@@ -35,7 +41,7 @@ embig l = concat $ take 5 $ iterate (map (map succW)) [concat $ take 5 (iterate
  4
  5      for each vertex v in Graph:
  6          dist[v] ← INFINITY
- 7          prev[v] ← UNDEFINED
+ 7          # prev[v] ← UNDEFINED
  8          add v to Q
  9      dist[source] ← 0
 10
@@ -48,19 +54,19 @@ embig l = concat $ take 5 $ iterate (map (map succW)) [concat $ take 5 (iterate
 17              alt ← dist[u] + length(u, v)
 18              if alt < dist[v]:
 19                  dist[v] ← alt
-20                  prev[v] ← u
+20                  # prev[v] ← u
 21
-22      return dist[], prev[]
+22      return dist[]
 -}
 
 type PQueue a = IntMap [a]
 
-pminView :: PQueue a -> (a, PQueue a)
+-- pminView :: PQueue a -> ((Key, a), PQueue a)
 pminView p =
   let Just (l, p') = IM.minViewWithKey p
    in case l of
         (_, []) -> pminView p
-        (k, x : xs) -> (x, if null xs then IM.delete k p' else IM.insert k xs p')
+        (k, x : xs) -> ((k, x), if null xs then IM.delete k p' else IM.insert k xs p')
 
 pins :: Int -> a -> PQueue a -> PQueue a
 pins k x = IM.insertWith (++) k [x]
@@ -73,34 +79,36 @@ neighbors g p = IS.filter (`IM.member` g) (IS.fromList (conv <$> [(x -1, y), (x,
   where
     (x, y) = from p
 
-dijkstra :: IntMap Int -> Int -> (IntMap Int, IntMap Int)
-dijkstra g s = step1 & step2 & step3
+dijkstra :: Grid -> Int -> IntMap Int
+dijkstra g s = step1 & step2
   where
-    step1 = let c = IM.keysSet g; (a, b) = IS.foldl' go (mempty, mempty) c in (a, b, c)
-    go (dv, pv) v = (dv', pv')
+    step1 = let c = IM.keysSet g; a = IS.foldl' go mempty c in (a, c, pins 0 s pempty)
       where
-        dv' = IM.insert v (maxBound :: Int) dv
-        pv' = IM.insert v (conv (0, 0)) pv
-    step2 (dv, pv, q) = (IM.insert s (conv (0, 0)) dv, pv, q)
-    step3 (dv, pv, q)
-      | IS.null q = (dv, pv)
-      | otherwise = step3 (dv', pv', q')
+        go dv v = IM.insert v (maxBound :: Int) dv
+    step2 (dv, unseen, q)
+      | IM.null q = dv
+      | otherwise = step2 (dv', unseen', q'')
       where
-        q' = IS.delete u q
-        u = nextVert dv
-        (dv', pv') = IS.foldl' f (dv, pv) (IS.intersection q (neighbors g u))
+        -- u is the closest vertex to visit
+        ((d, u), q') = pminView q
+        unseen' = IS.delete u unseen
+        -- for all neighbors of u, visit it and thread the distance map and priority queue
+        (dv', q'') = IS.foldl' visit (dv, q') (neighbors g u)
           where
-            f (dv, pv) v = (dv'', pv'')
+            -- to visit a point p', make sure it's in the graph
+            visit (dv, q') p'
+              | p' `IM.member` g = (dv'', q'')
+              | otherwise = (dv, q')
               where
-                alt = dv IM.! u + g IM.! v
-                (dv'', pv'') = if alt < dv IM.! v then (IM.insert v alt dv, IM.insert v u pv) else (dv, pv)
-        nextVert :: IntMap Int -> Int
-        nextVert dv = v
-          where
-            v :: Int
-            v = fst $ IS.foldl' (\(x, c) y -> if (dv IM.! y) < c then (y, dv IM.! y) else (x, c)) (undefined, maxBound :: Int) q
+                -- compute the new distance to p' (store into alt) and
+                -- compare against the previous distance
+                alt = d + g IM.! p'
+                -- if it's better then update distance for p' in dv to be alt
+                (dv'', q'') = if alt < dv IM.! p' then (IM.insert p' alt dv, pins alt p' q') else (dv, q')
 
-part1 (inp, n) = fst (dijkstra inp (conv (1, 1))) IM.! conv (n, n)
+part1 (inp, n) = dijkstra inp (conv (1, 1)) IM.! conv (n, n)
+
+part2 = part1
 
 main = do
   let dayNumber = 15 :: Int
@@ -110,14 +118,13 @@ main = do
   let x = map (map (read . pure)) . lines $ inp'
   let inp = toGrid x
   let inp2 = toGrid2 x
-  let n = length (head $ lines $ inp')
+  let n = length (head $ lines inp')
   print (part1 (inp, n))
   print (part1 (inp2, n * 5))
-
--- -- defaultMain
--- --   [ bgroup
--- --       dayString
--- --       [ bench "part1" $ whnf part1 inp,
--- --         bench "part2" $ whnf part2 inp
--- --       ]
--- --   ]
+  defaultMain
+    [ bgroup
+        dayString
+        [ bench "part1" $ whnf part1 (inp, n),
+          bench "part2" $ whnf part2 (inp2, n * 5)
+        ]
+    ]

readme.md

@@ -348,6 +348,26 @@ variance introduced by outliers: 55% (severely inflated)
 ```
 </details>
 
+### Day 15
+<details>
+
+```
+benchmarking day15/part1
+time                 28.84 ms   (28.27 ms .. 29.51 ms)
+                     0.997 R²   (0.992 R² .. 0.999 R²)
+mean                 29.88 ms   (29.11 ms .. 31.74 ms)
+std dev              2.642 ms   (879.3 μs .. 4.691 ms)
+variance introduced by outliers: 34% (moderately inflated)
+
+benchmarking day15/part2
+time                 1.172 s    (1.088 s .. 1.259 s)
+                     0.999 R²   (0.997 R² .. 1.000 R²)
+mean                 1.213 s    (1.190 s .. 1.244 s)
+std dev              29.96 ms   (12.16 ms .. 39.32 ms)
+variance introduced by outliers: 19% (moderately inflated)
+```
+</details>
+
 ### Day 16
 <details>