10-1 with dynamic programming / memoization

1 files changed, 9 insertions, 5 deletions

diff --git a/10-1.hs b/10-1.hs
index afeff45..c11c017 100644
--- a/10-1.hs
+++ b/10-1.hs
@@ -1,3 +1,5 @@
+-- compile with -O2 to take advantage of memoization!
+
 import Data.Array
 import Data.List (nub)
 
@@ -5,6 +7,11 @@ data Pos = Pos Int Int deriving (Eq,Ix,Ord,Show)
 
 type Map = Array Pos Int
 
+type Memo = Array Pos [Pos]
+
+memo :: Map -> Memo
+memo m = array (bounds m) [ (p, reachable m p) | p <- indices m ]
+
 step :: Map -> Pos -> [Pos]
 step m p@(Pos x y) = filter ((== 1) . subtract (m ! p) . (m !))
   . filter (`elem` indices m)
@@ -22,10 +29,7 @@ height = length
 
 reachable :: Map -> Pos -> [Pos]
 reachable m p | m ! p == 9 = [p]
-              | otherwise  = concatMap (reachable m) $ step m p
-
-score :: Map -> Pos -> Int
-score m = length . nub . reachable m
+              | otherwise  = nub . concatMap (memo m !) $ step m p
 
 trailheads :: Map -> [Pos]
 trailheads = map fst . filter ((== 0) . snd) . assocs
@@ -38,4 +42,4 @@ parse ls = listArray
 main :: IO ()
 main = do
   m <- parse . lines <$> getContents
-  print . sum . map (score m) . trailheads $ m
+  print . sum . map (length . (memo m !)) . trailheads $ m