day14 rewrite + part2

2023/day14/day14.hs

@@ -1,50 +1,39 @@
 import Control.Arrow
-import Control.Lens
+import Data.List
+import Data.List.Split
 import qualified Data.Map as M
-import Linear.V2
-import Safe
-
-type Grid = M.Map (V2 Int) Char
 
 main :: IO ()
-main = interact $ show . day14 . parse . lines
+main = interact $ show . (day14a &&& day14b 1000000000) . lines
 
-day14 :: Grid -> Int
-day14 m =
-  sum $ map ((maxY -) . view _y . fst) (M.toList $ M.filter (== 'O') (moveAll m))
-      where maxY = succ $ maximum $ map (view _y) $ M.keys m
+day14a :: [String] -> Int
+day14a = calc . turn . turn . turn . gravity . turn
 
-moveAll :: Grid -> Grid
-moveAll m = foldl moveUp m os
+turn :: [String] -> [String]
+turn = transpose . reverse
+
+gravity :: [String] -> [String]
+gravity = map (intercalate "#" . map sort . splitOn "#")
+
+calc :: [String] -> Int
+calc = sum . zipWith (*) [1 ..] . map (count 'O') . reverse
   where
-    os = map fst $ M.toList $ M.filter (== 'O') m
+    count c = length . filter (== c)
 
-parse :: [String] -> Grid
-parse = M.fromList . concat . zipWith (map . f) [0 ..] . map (zip [0 ..])
+genStates :: [String] -> [[String]]
+genStates = map snd . filter f . zip [0 ..] . iterate (gravity . turn)
   where
-    f x (y, z) = (V2 y x, z)
+    f (x, _) = x `mod` 4 == 0
 
-unParse :: Grid -> String
-unParse m =
-  concat
-    [ M.elems (M.filterWithKey (\(V2 _ x) _ -> x == n) m) ++ "\n"
-    | n <- [a .. b]
-    ]
+findDupe :: M.Map [String] Int -> Int -> [[String]] -> (Int, Int)
+findDupe m n (x:xs) =
+  if x `M.member` m
+    then (m M.! x, n)
+    else findDupe (M.insert x n m) (succ n) xs
+
+day14b :: Int -> [String] -> Int
+day14b n input = calc (states !! num)
   where
-    (a, b) = (minimum &&& maximum) $ map (view _y) $ M.keys m
-
-findLetter :: Char -> Grid -> V2 Int
-findLetter c = fst . head . filter ((== c) . snd) . M.toList
-
-filterLowestDot :: V2 Int -> V2 Int -> Char -> Bool
-filterLowestDot (V2 x1 y1) (V2 x2 y2) c = c /= 'O' && x1 == x2 && y1 > y2
-
-moveUp :: Grid -> V2 Int -> Grid
-moveUp m k = swapCoord m newCoord k
-  where
-    line = M.toList $ M.filterWithKey (filterLowestDot k) m
-    newCoord = lastMay $ takeWhile (\x -> snd x /= '#') $ reverse line
-
-swapCoord :: Grid -> Maybe (V2 Int, Char) -> V2 Int -> Grid
-swapCoord m Nothing _ = m
-swapCoord m (Just (v2, c)) v1 = M.insert v1 c $ M.insert v2 (m M.! v1) m
+    num = (n - s1) `mod` (s2 - s1) + s1
+    (s1, s2) = findDupe M.empty 0 states
+    states = genStates input