Day11 part 2 and refactoring

src/Lib/Matrix.hs

@@ -1,4 +1,4 @@
-module Lib.Matrix (Matrix, (!), (!?), height, width, insertCol, insertRow, getRow, getCol, getCols, findColIndices, findRowIndices, Lib.Matrix.toList, findLocations) where
+module Lib.Matrix (Matrix, (!), (!?), height, width, insertCol, insertRow, getRow, getCol, getCols, findColIndices, findRowIndices, Lib.Matrix.toList, findLocations, Lib.Matrix.map, updateRow, updateCol) where
 
 import Data.Foldable (toList)
 import Data.List (findIndices)
@@ -25,9 +25,15 @@ width = S.length . (`S.index` 0)
 insertCol :: Seq a -> Int -> Matrix a -> Matrix a
 insertCol col colIdx = fmap (S.insertAt colIdx (col `S.index` colIdx))
 
+updateCol :: Seq a -> Int -> Matrix a -> Matrix a
+updateCol col colIdx = fmap (S.update colIdx (col `S.index` colIdx))
+
 insertRow :: Seq a -> Int -> Matrix a -> Matrix a
 insertRow row rowIdx = S.insertAt rowIdx row
 
+updateRow :: Seq a -> Int -> Matrix a -> Matrix a
+updateRow row rowIdx = S.update rowIdx row
+
 getRow :: Int -> Matrix a -> Seq a
 getRow = flip S.index
 
@@ -48,3 +54,6 @@ toList mat = Data.Foldable.toList $ fmap Data.Foldable.toList mat
 
 findLocations :: (a -> Bool) -> Matrix a -> [(Int, Int)]
 findLocations predicate mat = [(x, y) | y <- [0 .. height mat - 1], x <- [0 .. width mat - 1], predicate (mat ! (x, y))]
+
+map :: (a -> b) -> Matrix a -> Matrix b
+map f = fmap (fmap f)

src/Solutions/Day11.hs

@@ -1,82 +1,47 @@
 {-# LANGUAGE TypeFamilies #-}
 
-module Solutions.Day11 (day11, parser, part1, part2) where
+module Solutions.Day11 (day11, parser, part1, part2, sumOfGalaxyDistances) where
 
-import Data.Array (Array)
-import qualified Data.Array as A
 import qualified Data.Sequence as S
 import Lib.AOC (runSolution)
-import Lib.SearchAlgorithms (DijkstraGraph(..), Distance(..), findShortestDistance)
 import Lib.Matrix (Matrix)
 import qualified Lib.Matrix as Mat
-import Data.Foldable (foldl', toList)
-import Debug.Trace
-import Data.Sequence (Seq((:<|)))
-import Data.Maybe (catMaybes)
-import Lib.List (pairs)
 
 type Input = Matrix Char
 
 parser :: String -> Input
-parser =  S.fromList . map S.fromList . lines
-
-isSpace :: Char -> Bool
-isSpace = (==) '.'
-
--- Modifes the indexes to insert at to account for the fact that insert one index will change the position where
--- the next insertion needs to occur.
-indicesToInsertAt :: [Int] -> [Int]
-indicesToInsertAt [] = []
-indicesToInsertAt (x:xs) = x : indicesToInsertAt (map (+ 1) xs)
-
-expandUniverse :: Input -> Input
-expandUniverse universe =
-  -- We need to use indicesToInsertAt because each time we add a row the indexes where we need to insert new rows will move down
-  -- by 1. Same logic applies to inserting columns.
-  let emptyRowIndices =  indicesToInsertAt $ Mat.findRowIndices (all isSpace) universe
-      emptyRow = S.fromList $ replicate (Mat.width universe) '.'
-      universeWithRowsExpanded = foldl' (flip $ Mat.insertRow emptyRow) universe emptyRowIndices
-      emptyColIndices = indicesToInsertAt $ Mat.findColIndices (all isSpace) universeWithRowsExpanded
-      emptyCol = S.fromList $ replicate (Mat.height universeWithRowsExpanded) '.'
-  in foldl' (flip $ Mat.insertCol emptyCol) universeWithRowsExpanded emptyColIndices
-
--- Just for debugging the state of the "universe" to see what it looks like after expansion
-prettyPrint :: Input -> String
-prettyPrint S.Empty = ""
-prettyPrint (x :<| xs) = show (toList x) ++ "\n" ++ prettyPrint xs
-
-part1 :: Input -> Int
-part1 input = distances galaxyLocations
+parser = S.fromList . map S.fromList . lines
+
+getExpandedGalaxyLocations :: Int -> Input -> [(Int, Int)]
+getExpandedGalaxyLocations expansionFactor universe = 
+  [ expandedLoc | 
+    (locX, locY) <- Mat.findLocations (== '#') universe,
+        let 
+          xIncrease = (expansionFactor - 1) * length (filter (< locX) emptyCols)
+          yIncrease = (expansionFactor - 1) * length (filter (< locY) emptyRows)
+          expandedLoc = (locX + xIncrease, locY + yIncrease)
+      ]
   where
-    expandedUniverse = expandUniverse input
-    galaxyLocations = Mat.findLocations (== '#') expandedUniverse
-    findDistance (x0, y0) (x1, y1) = abs (x1 - x0) + abs (y1 - y0)
+    emptyRows = Mat.findRowIndices (all (== '.')) universe
+    emptyCols = Mat.findColIndices (all (== '.')) universe
 
-    distances [] = 0
-    distances (x:xs) = sum (map (findDistance x) xs) + distances xs
+sumOfGalaxyDistances :: Int -> Input -> Int
+sumOfGalaxyDistances expansionFactor universe = distances galaxyLocations
+ where
+  galaxyLocations = getExpandedGalaxyLocations expansionFactor universe
 
-newtype Graph2D = Graph2D (Array (Int, Int) Char) deriving (Show)
+  -- Straight line distance between two points
+  findDistance (x0, y0) (x1, y1) = abs (x0 - x1) + abs (y0 - y1)
 
--- matrixToGraph2D :: Input -> Graph2D
--- matrixToGraph2D mat = Graph2D $ A.listArray ((0, 0), (Mat.height mat - 1, Mat.width mat - 1)) $ concat $ Mat.toList mat
+  distances [] = 0
+  distances (x : xs) = sum (map (findDistance x) xs) + distances xs
 
--- instance DijkstraGraph Graph2D where
---   type DijkstraNode Graph2D = (Int, Int)
---   type DijkstraCost Graph2D = Int
---   -- the cost is always 1 for our graph
---   dijkstraEdges (Graph2D arr) node = [(neighbor, 1) | neighbor <- getNeighbors node arr]
 
--- getNeighbors :: (Int, Int) -> Array (Int, Int) Char -> [(Int, Int)]
--- getNeighbors (x, y) arr = catMaybes [up, down, left, right]
---   where
---     (maxRow, maxCol) = snd $ A.bounds arr
---     up = if y - 1 < 0 then Nothing else Just (x, y - 1)
---     down = if y + 1 > maxRow then Nothing else Just (x, y + 1)
---     left = if x - 1 < 0 then Nothing else Just (x - 1, y)
---     right = if x + 1 > maxCol then Nothing else Just (x + 1, y)
+part1 :: Input -> Int
+part1 = sumOfGalaxyDistances 2
 
-part2 :: Input -> String
-part2 = show . const 0
+part2 :: Input -> Int
+part2 = sumOfGalaxyDistances 1_000_000
 
 day11 :: IO ()
 day11 = runSolution "11" parser part1 part2

test/Day11.hs

@@ -3,7 +3,7 @@
 module Day11 where
 
 import qualified Data.Sequence as S
-import Solutions.Day11 (parser, part1, part2)
+import Solutions.Day11 (parser, part1, part2, sumOfGalaxyDistances)
 import Test.Tasty
 import Test.Tasty.HUnit
 
@@ -45,10 +45,13 @@ test_day11 =
         input <- parser <$> readFile "data/day11.txt"
         part1 input @?= 9445168
 
-    -- , testCase "part 2 - example input" $ do
-    --     part2 (parser exampleInput) @?= 2
+    , testCase "part 2 - example input" $ do
+        sumOfGalaxyDistances 10 (parser exampleInput) @?= 1030
+
+    , testCase "part 2 - example input (factor 100)" $ do
+        sumOfGalaxyDistances 100 (parser exampleInput) @?= 8410 
 
     -- , testCase "part 2" $ do
     --     input <- parser <$> readFile "data/day11.txt"
-    --     part2 input @?= 957
+    --     part2 input @?= 100 -- TODO: get real value
     ]