src/advent22/advent22a.hs

   1 {-# LANGUAGE NegativeLiterals #-}
   2 {-# LANGUAGE FlexibleContexts #-}
   3 {-# LANGUAGE OverloadedStrings #-}
   4 {-# LANGUAGE TypeFamilies #-}
   5 {-# LANGUAGE BangPatterns #-}
   6
   7 import Prelude hiding (Left, Right)
   8 import Data.List
   9 import qualified Data.Set as S
  10
  11 type Point = (Int, Int)
  12 type Infection = S.Set Point
  13
  14 data Direction = Up | Right | Down | Left deriving (Show, Eq, Enum)
  15
  16 leftOf Up = Left
  17 leftOf x = pred x
  18
  19 rightOf Left = Up
  20 rightOf x = succ x
  21
  22 delta :: Direction -> Point
  23 delta Up = (-1, 0)
  24 delta Right = (0, 1)
  25 delta Down = (1, 0)
  26 delta Left = (0, -1)
  27
  28 (+:) :: Point -> Point -> Point
  29 (+:) (r, c) (dr, dc) = (r + dr, c + dc)
  30
  31 data World = World { infected :: Infection
  32                    , position :: Point
  33                    , direction :: Direction
  34                    , infectionCount :: Int
  35                    } deriving (Eq, Show)
  36
  37
  38 main :: IO ()
  39 main = do
  40         text <- readFile "data/advent22.txt"
  41         let grid = lines text
  42         print $ infectionCount $ progress 10000 $ initialWorld grid
  43
  44 initialWorld :: [String] -> World
  45 initialWorld grid = World
  46     { infected = initialInfected grid
  47     , position = initialPosition grid
  48     , direction = Up
  49     , infectionCount = 0
  50     }
  51
  52 initialInfected :: [String] -> Infection
  53 initialInfected g = S.fromList [(r, c) | r <- [0..(length g - 1)]
  54                                        , c <- [0..((length . head) g - 1)]
  55                                        , g!!r!!c == '#']
  56
  57 initialPosition :: [String] -> Point
  58 initialPosition g = (length g `div` 2, (length . head) g `div` 2)
  59
  60
  61 progress :: Int -> World -> World
  62 progress n = (!! n) . iterate step
  63
  64 step :: World -> World
  65 step world = World { infected = inf', position = pos', direction = dir'
  66                    , infectionCount = ic'}
  67     where here = position world
  68           infectedHere = here `S.member` infected world
  69           dir' = if infectedHere then rightOf (direction world)
  70                                  else leftOf (direction world)
  71           inf' = if infectedHere then S.delete here $ infected world
  72                                  else S.insert here $ infected world
  73           ic'  = if infectedHere then infectionCount world
  74                                  else infectionCount world + 1
  75           pos' = here +: delta dir'