full memoization

2021/day12.hs

@@ -1,41 +1,78 @@
 module Day12 where
-import Data.Map ( Map, fromList, unionWith, empty, (!) )
+-- import Data.HashSet (HashSet, member, insert, fromList)
+-- import Data.Map ( Map, fromList, unionWith, empty, (!), insert, elems) qualified (map)
+import qualified Data.Map as M
+import Data.Bits ( Bits(shiftL, (.|.), (.&.)) )
+
+import qualified Data.Array as A
+
+import Data.HashSet (HashSet)
+import qualified Data.HashSet as S
 import Data.List.Split (splitOn)
 import qualified Data.Char
 
-type Path = (String,String)
-type System = Map String [String]
+import Debug.Trace ( trace )
+import GHC.Base (VecElem(Int16ElemRep))
 
-readPath :: String -> Path
+type Cave = Int
+type Path = (Cave,Cave)
+type System = M.Map Cave [Cave]
+
+
+readPath :: String -> (String,String)
 readPath s = (from,to) where [from,to] = splitOn "-" s
 
-addPath :: Path -> System -> System
-addPath (from,to) = unionWith (++) (Data.Map.fromList [(from,[to]), (to,[from])])
+addPath :: (String,String) -> M.Map String [String] -> M.Map String [String]
+addPath (from,to) = M.unionWith (++) (M.fromList [(from,[to]), (to,[from])])
 
-isUpper :: String -> Bool
-isUpper = Data.Char.isUpper . head
+isUpper :: Cave -> Bool
+isUpper x = x < 0
 
-pathFinder :: System -> [String] -> [[String]]
-pathFinder sys p@("end":xs) = [p]
-pathFinder sys p@(x:xs) = concatMap (\next -> pathFinder sys (next:p)) elig
-    where elig = filter (\cave -> isUpper cave || cave `notElem` p) (sys ! x)
-
-pathFinder2 :: System -> [String] -> [[String]]
-pathFinder2 sys p@("end":xs) = [p]
-pathFinder2 sys p@(x:xs) = concatMap (\next -> f next sys (next:p)) elig
+solve :: System -> (Int, Int)
+solve sys = (pf1a A.! (1,1), pf2a A.! (1, 1))
     where
-        f cave = if isUpper cave || cave `notElem` p then pathFinder2 else pathFinder
-        elig = filter (/= "start") (sys ! x)
+        caves = concat $ M.elems sys
+        bounds = ((minimum caves, 0), (maximum caves, maximum caves * 2))
+
+        pf1a :: A.Array (Cave,Int) Int
+        pf1a = A.array bounds [((c,v), pf1  c v) | c <- [minimum caves..maximum caves], v <- [0..maximum caves * 2] ]
+        pf1 :: Cave -> Int -> Int
+        pf1 (-1) _ = 1
+        pf1 x visited = sum $ map (\next -> pf1a A.! (next,nextmap next)) elig
+            where
+                nextmap next = if not (isUpper next) then next .|. visited else visited
+                elig = filter (\cave -> isUpper cave || (cave .&. visited) == 0) (sys M.! x)
+
+        pf2a :: A.Array (Cave,Int) Int
+        pf2a = A.array bounds [((c,v), pf2 c v) | c <- [minimum caves..maximum caves], v <- [0..maximum caves * 2] ]
+        pf2 :: Cave -> Int -> Int
+        pf2 (-1) _ = 1
+        pf2 x visited = sum $ map (\next -> f next A.! (next,nextmap next)) elig
+            where
+                nextmap next = if not (isUpper next) then next .|. visited else visited
+                f cave = if isUpper cave || (cave .&. visited) == 0 then pf2a else pf1a
+                elig = filter (/= 1) (sys M.! x)
+
+convert :: [String] -> M.Map String Int
+convert caves = conversion
+    where
+        (conversion, numhi, numlo) = foldr go (M.fromList [("end", -1), ("start", 1)], 2, 1) caves
+        go cave (map,upper,lower)
+          | cave `M.member` map           = (map,upper,lower)
+          | Data.Char.isUpper (head cave) = (M.insert cave (-upper) map, upper+1, lower)
+          | otherwise                     = (M.insert cave (1 `Data.Bits.shiftL` lower) map, upper, lower+1)
 
 main :: IO ()
 main = do
     input <- map readPath . lines <$> getContents
-    let system = foldr addPath Data.Map.empty input
+    let system = foldr addPath M.empty input
+    let conversion = convert $ concat $ M.elems system
+    let integralSystem = M.mapKeys (conversion M.!) $ M.map (map (conversion M.!)) system
+
+    let (p1,p2) = solve integralSystem
 
     putStrLn "part 1: "
-    let paths = map reverse $ pathFinder system ["start"]
-    print $ length paths
+    print p1
 
     putStrLn "part 2: "
-    let paths = map reverse $ pathFinder2 system ["start"] 
-    print $ length paths
+    print p2