src/advent22/advent22b.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.Map.Strict as M
  10
  11 type Point = (Int, Int)
  12
  13 data Flag = Clean | Weakened | Infected | Flagged deriving (Show, Eq)
  14
  15 type Infection = M.Map Point Flag
  16
  17 data Direction = Up | Right | Down | Left deriving (Show, Eq, Enum)
  18
  19 leftOf Up = Left
  20 leftOf x = pred x
  21
  22 rightOf Left = Up
  23 rightOf x = succ x
  24
  25 delta :: Direction -> Point
  26 delta Up = (-1, 0)
  27 delta Right = (0, 1)
  28 delta Down = (1, 0)
  29 delta Left = (0, -1)
  30
  31 (+:) :: Point -> Point -> Point
  32 (+:) (r, c) (dr, dc) = (r + dr, c + dc)
  33
  34 data World = World { infected :: Infection
  35                    , position :: Point
  36                    , direction :: Direction
  37                    , infectionCount :: Int
  38                    } deriving (Eq, Show)
  39
  40
  41 main :: IO ()
  42 main = do
  43         text <- readFile "data/advent22.txt"
  44         let grid = lines text
  45         print $ infectionCount $ progress 10000000 $ initialWorld grid
  46
  47 initialWorld :: [String] -> World
  48 initialWorld grid = World
  49     { infected = initialInfected grid
  50     , position = initialPosition grid
  51     , direction = Up
  52     , infectionCount = 0
  53     }
  54
  55 initialInfected :: [String] -> Infection
  56 initialInfected g = M.fromList [ ((r, c), Infected)
  57                                | r <- [0..(length g - 1)]
  58                                , c <- [0..((length . head) g - 1)]
  59                                , g!!r!!c == '#']
  60
  61 initialPosition :: [String] -> Point
  62 initialPosition g = (length g `div` 2, (length . head) g `div` 2)
  63
  64
  65 progress :: Int -> World -> World
  66 progress n = (!! n) . iterate step
  67
  68
  69 step :: World -> World
  70 step world = World { infected = inf', position = pos', direction = dir'
  71                    , infectionCount = ic'}
  72     where !here = position world
  73           !stateHere = M.findWithDefault Clean here (infected world)
  74           !dir' = case stateHere of
  75                       Clean -> leftOf (direction world)
  76                       Weakened -> direction world
  77                       Infected -> rightOf (direction world)
  78                       Flagged -> rightOf (rightOf (direction world))
  79           !stateHere' = case stateHere of
  80                       Clean -> Weakened
  81                       Weakened -> Infected
  82                       Infected -> Flagged
  83                       Flagged -> Clean
  84           !inf' = M.insert here stateHere' (infected world)
  85           !ic'  = if stateHere' == Infected then infectionCount world + 1
  86                                            else infectionCount world
  87           !pos' = here +: delta dir'