1 files changed, 86 insertions, 0 deletions
diff --git a/2020/11.hs b/2020/11.hs
new file mode 100644
index 0000000..2d573be
--- /dev/null
+++ b/2020/11.hs
@@ -0,0 +1,86 @@
+{-# LANGUAGE LambdaCase #-}
+module Main where
+
+import Data.List (find, findIndex, intercalate)
+import Data.Maybe (catMaybes, fromJust)
+import qualified Data.Vector as V
+import qualified Data.Vector.Storable as U
+import Data.Word (Word8)
+import Foreign.Ptr (Ptr, castPtr)
+import Foreign.Storable (Storable(..))
+
+import Input
+
+
+data Cell = Floor | Seat | Occupied
+  deriving (Show, Eq, Enum)
+
+instance Storable Cell where
+    sizeOf _ = sizeOf (undefined :: Word8)
+    alignment _ = alignment (undefined :: Word8)
+    peek p = toEnum . fromIntegral <$> peek (castPtr p :: Ptr Word8)
+    poke p x = poke (castPtr p :: Ptr Word8) (fromIntegral (fromEnum x))
+
+parseCell :: Char -> Cell
+parseCell '.' = Floor
+parseCell 'L' = Seat
+parseCell '#' = Occupied
+parseCell _ = error "Cannot parse cell"
+
+showCell :: Cell -> Char
+showCell Floor = '.'
+showCell Seat = 'L'
+showCell Occupied = '#'
+
+data Grid = Grid Int Int (U.Vector Cell)
+  deriving (Show, Eq)
+
+parseGrid :: [String] -> Grid
+parseGrid lns =
+    let w = length (head lns)
+        h = length lns
+    in Grid w h (U.fromListN (w * h) (concatMap (map parseCell) lns))
+
+showGrid :: Grid -> String
+showGrid (Grid w h v) =
+    intercalate "\n" [[showCell (v U.! (w * y + x)) | x <- [0..w-1]]
+                     | y <- [0..h-1]]
+
+numOccupied :: Grid -> Int
+numOccupied (Grid _ _ v) = sum (map (fromEnum . (== Occupied)) (U.toList v))
+
+evolve :: Grid -> ([Cell] -> Maybe Int) -> Int -> Grid -> Grid
+evolve (Grid w h template) nbFunc scareCount =
+    \(Grid _ _ v) ->
+        let newCellAt x y =
+                let neighbours = [v U.! (w * y' + x') | (x', y') <- nbIndicesArr V.! (w * y + x)]
+                    filledNbs = sum (map (fromEnum . (== Occupied)) neighbours)
+                in case v U.! (w * y + x) of
+                     Floor -> Floor
+                     Seat | filledNbs == 0 -> Occupied
+                          | otherwise -> Seat
+                     Occupied | filledNbs >= scareCount -> Seat
+                              | otherwise -> Occupied
+        in Grid w h (U.generate (w * h) (\i -> newCellAt (i `rem` w) (i `quot` w)))
+  where
+    inBounds (x, y) = 0 <= x && x < w && 0 <= y && y < h
+    nbIndices x y = catMaybes
+                      [(sight !!) <$> nbFunc [template U.! (w * y' + x') | (x', y') <- sight]
+                      | dy <- [-1..1], dx <- [-1..1]
+                      , dx /= 0 || dy /= 0
+                      , let sight = takeWhile inBounds [(x+i*dx, y+i*dy) | i <- [1..]]]
+    nbIndicesArr = V.generate (w * h) (\i -> nbIndices (i `rem` w) (i `quot` w))
+
+
+equilibriumFor :: Eq a => (a -> a) -> a -> a
+equilibriumFor f x = fst (fromJust (find (uncurry (==)) (zip xs (tail xs))))
+  where xs = iterate f x
+
+main :: IO ()
+main = do
+    input <- getInput 11
+    let initGrid = parseGrid input
+        evolve1 = evolve initGrid (\case Seat:_ -> Just 0 ; _ -> Nothing) 4
+        evolve2 = evolve initGrid (findIndex (== Seat)) 5
+    print (numOccupied (equilibriumFor evolve1 initGrid))
+    print (numOccupied (equilibriumFor evolve2 initGrid))