10-2 with memoization

1 files changed, 23 insertions, 4 deletions

diff --git a/10-2.hs b/10-2.hs
index cf2840c..6ee5bc2 100644
--- a/10-2.hs
+++ b/10-2.hs
@@ -1,8 +1,11 @@
+import Control.Monad.State
 import Data.Array
+import Data.Foldable (toList)
 
 data Pos = Pos Int Int deriving (Eq,Ix,Ord,Show)
 
 type Map = Array Pos Int
+type S = State (Array Pos (Maybe Int)) Int
 
 step :: Map -> Pos -> [Pos]
 step m p@(Pos x y) = filter ((== 1) . subtract (m ! p) . (m !))
@@ -19,9 +22,19 @@ width,height :: [[a]] -> Int
 width = length . head
 height = length
 
-score :: Map -> Pos -> Int
-score m p | m ! p == 9 = 1
-          | otherwise  = sum . map (score m) $ step m p
+score :: Map -> Pos -> S
+score m p | m ! p == 9 = return 1
+          | otherwise  = fmap sum . traverse (score' m) $ step m p
+
+score' :: Map -> Pos -> S
+score' m p = do
+  memo <- get
+  case memo ! p of
+    Just x  -> return x
+    Nothing -> do
+      x <- score m p
+      modify (// [(p,Just x)])
+      return x
 
 trailheads :: Map -> [Pos]
 trailheads = map fst . filter ((== 0) . snd) . assocs
@@ -34,4 +47,10 @@ parse ls = listArray
 main :: IO ()
 main = do
   m <- parse . lines <$> getContents
-  print . sum . map (score m) . trailheads $ m
+  print
+    . sum
+    . toList
+    . flip evalState (listArray (bounds m) $ repeat Nothing)
+    . traverse (score' m)
+    . trailheads
+    $ m