advent11/src/advent11.hs

   1 -- import Debug.Trace
   2
   3 import Prelude hiding (Left, Right)
   4
   5 import qualified Data.Map as M
   6 import Data.Map ((!))
   7 import qualified Data.Set as S
   8 -- import Data.Sort
   9 import Data.List
  10 import Control.Monad.Reader
  11 import Control.Monad.Loops
  12
  13
  14 type Position = (Int, Int)
  15 data Seat = Floor | Empty | Occupied deriving (Eq, Ord)
  16 data Direction = Up | UpRight | Right | DownRight | Down | DownLeft | Left | UpLeft
  17   deriving (Eq, Ord, Show, Enum)
  18 type Seats = M.Map Position Seat
  19
  20 type Neighbourhood = M.Map Position (S.Set Position)
  21 type Rule = Seats -> Neighbourhood -> Position -> Seat -> Seat
  22
  23 instance Show Seat where
  24   show Floor = "."
  25   show Empty = "L"
  26   show Occupied = "#"
  27
  28
  29 type CachedSeats a = Reader (Neighbourhood, Rule) a
  30
  31
  32 main :: IO ()
  33 main =
  34   do  text <- readFile "data/advent11.txt"
  35       let (seats, maxCorner) = readGrid text
  36       -- print $ M.size seats
  37       -- print maxCorner
  38       print $ part1 seats
  39       print $ part2 seats
  40
  41
  42 part1 seats = M.size $ M.filter (== Occupied) stableSeats
  43   where cachedNeighbours = allNeighbourhoods seats
  44         env = (cachedNeighbours, ruleA)
  45         stableSeats = snd $ runReader (runSteps seats) env
  46
  47 part2 seats = M.size $ M.filter (== Occupied) stableSeats
  48   where cachedNeighbours = allSightNeighbourhoods seats
  49         env = (cachedNeighbours, ruleB)
  50         stableSeats = snd $ runReader (runSteps seats) env
  51
  52
  53 runSteps :: Seats -> CachedSeats (Seats, Seats)
  54 runSteps seats = iterateUntilM (uncurry (==)) seatChanges (M.empty, seats)
  55
  56 seatChanges :: (Seats, Seats) -> CachedSeats (Seats, Seats)
  57 seatChanges (_, seats0) =
  58   do seats <- step seats0
  59      return (seats0, seats)
  60
  61 step :: Seats -> CachedSeats Seats
  62 step seats =
  63   do  (nbrs, rule) <- ask
  64       return $ M.mapWithKey (rule seats nbrs) seats
  65
  66 ruleA :: Seats -> Neighbourhood -> Position -> Seat -> Seat
  67 ruleA seats nbrs here thisSeat
  68   | thisSeat == Empty && nOccs == 0 = Occupied
  69   | thisSeat == Occupied && nOccs >= 4 = Empty
  70   | otherwise = thisSeat
  71   where nOccs = M.size $ occupiedNeighbours seats nbrs here
  72
  73 ruleB :: Seats -> Neighbourhood -> Position -> Seat -> Seat
  74 ruleB seats nbrs here thisSeat
  75   | thisSeat == Empty && nOccs == 0 = Occupied
  76   | thisSeat == Occupied && nOccs >= 5 = Empty
  77   | otherwise = thisSeat
  78   where nOccs = M.size $ occupiedNeighbours seats nbrs here
  79
  80
  81 neighbours (r, c) = S.delete (r, c) $ S.fromList [(r + dr, c + dc) | dr <- [-1, 0, 1], dc <- [-1, 0, 1]]
  82
  83 neighbourhood seats here = S.intersection (M.keysSet seats) (neighbours here)
  84
  85 allNeighbourhoods :: Seats -> Neighbourhood
  86 allNeighbourhoods seats = M.mapWithKey (\h _ -> neighbourhood seats h) seats
  87
  88 occupiedNeighbours seats nbrs here = M.filter (== Occupied) $ M.restrictKeys seats (nbrs!here)
  89
  90
  91 onSightLine :: Position -> Direction -> Position -> Bool
  92 onSightLine (r0, c0) Down      (r, c) = (c0 == c) && (r > r0)
  93 onSightLine (r0, c0) Up        (r, c) = (c0 == c) && (r < r0)
  94 onSightLine (r0, c0) Right     (r, c) = (r0 == r) && (c > c0)
  95 onSightLine (r0, c0) Left      (r, c) = (r0 == r) && (c < c0)
  96 onSightLine (r0, c0) DownRight (r, c) = ((r - r0) > 0) && ((r - r0) == (c - c0))
  97 onSightLine (r0, c0) UpLeft    (r, c) = ((r - r0) < 0) && ((r - r0) == (c - c0))
  98 onSightLine (r0, c0) DownLeft  (r, c) = ((r - r0) > 0) && ((r - r0) == (c0 - c))
  99 onSightLine (r0, c0) UpRight   (r, c) = ((r - r0) < 0) && ((r - r0) == (c0 - c))
 100
 101 manhattan (r1, c1) (r2, c2) = abs (r1 - r2) + abs (c1 - c2)
 102
 103 closestInDirection seats here direction = take 1 sortedSeats
 104   -- where seatsInDirection = M.keys $ M.filterWithKey (\o _ -> onSightLine here direction o) seats
 105   where seatsInDirection = filter (onSightLine here direction) $ M.keys seats
 106         sortedSeats = sortOn (manhattan here) seatsInDirection
 107
 108 closestInSight :: Seats -> Position -> (S.Set Position)
 109 closestInSight seats here = S.fromList $ concatMap (closestInDirection seats here) [d | d <- [Up .. UpLeft]]
 110
 111 allSightNeighbourhoods :: Seats -> Neighbourhood
 112 allSightNeighbourhoods seats = M.mapWithKey (\h _ -> closestInSight seats h) seats
 113
 114 -- occupiedInSight :: Seats -> Position -> Seats
 115 -- occupiedInSight seats here = M.filter (== Occupied) $ M.restrictKeys seats $ closestInSight seats here
 116
 117
 118
 119 readGrid :: String -> (Seats, Position)
 120 readGrid input = (seats, (maxR, maxC))
 121   where seats = M.fromList $ concat
 122                   [ [((r, c), Empty) | (t, c) <- zip row [0..], t == 'L']
 123                   | (row, r) <- zip rows [0..] ]
 124         rows = lines input
 125         maxC = (length $ head rows) - 1
 126         maxR = (length rows) - 1
 127
 128 showGrid seats (maxR, maxC) =
 129   unlines $ [ concat [showSeat (r, c) | c <- [0..maxC] ] | r <- [0..maxR]]
 130   where showSeat here = show $ M.findWithDefault Floor here seats
 131