1 -- import Debug.Trace

3 import Prelude hiding (Left, Right)

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

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

20 type Neighbourhood = M.Map Position (S.Set Position)

21 type Rule = Seats -> Neighbourhood -> Position -> Seat -> Seat

23 instance Show Seat where

24 show Floor = "."

25 show Empty = "L"

26 show Occupied = "#"

29 type CachedSeats a = Reader (Neighbourhood, Rule) a

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

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

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

53 runSteps :: Seats -> CachedSeats (Seats, Seats)

54 runSteps seats = iterateUntilM (uncurry (==)) seatChanges (M.empty, seats)

56 seatChanges :: (Seats, Seats) -> CachedSeats (Seats, Seats)

57 seatChanges (_, seats0) =

58 do seats <- step seats0

59 return (seats0, seats)

61 step :: Seats -> CachedSeats Seats

62 step seats =

63 do (nbrs, rule) <- ask

64 return $ M.mapWithKey (rule seats nbrs) seats

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

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

81 neighbours (r, c) = S.delete (r, c) $ S.fromList [(r + dr, c + dc) | dr <- [-1, 0, 1], dc <- [-1, 0, 1]]

83 neighbourhood seats here = S.intersection (M.keysSet seats) (neighbours here)

85 allNeighbourhoods :: Seats -> Neighbourhood

86 allNeighbourhoods seats = M.mapWithKey (\h _ -> neighbourhood seats h) seats

88 occupiedNeighbours seats nbrs here = M.filter (== Occupied)

89 $ M.restrictKeys seats (nbrs!here)

92 onSightLine :: Position -> Direction -> Position -> Bool

93 onSightLine (r0, c0) Down (r, c) = (c0 == c) && (r > r0)

94 onSightLine (r0, c0) Up (r, c) = (c0 == c) && (r < r0)

95 onSightLine (r0, c0) Right (r, c) = (r0 == r) && (c > c0)

96 onSightLine (r0, c0) Left (r, c) = (r0 == r) && (c < c0)

97 onSightLine (r0, c0) DownRight (r, c) = ((r - r0) > 0) && ((r - r0) == (c - c0))

98 onSightLine (r0, c0) UpLeft (r, c) = ((r - r0) < 0) && ((r - r0) == (c - c0))

99 onSightLine (r0, c0) DownLeft (r, c) = ((r - r0) > 0) && ((r - r0) == (c0 - c))

100 onSightLine (r0, c0) UpRight (r, c) = ((r - r0) < 0) && ((r - r0) == (c0 - c))

102 manhattan (r1, c1) (r2, c2) = abs (r1 - r2) + abs (c1 - c2)

104 closestInDirection seats here direction = take 1 sortedSeats

105 -- where seatsInDirection = M.keys $ M.filterWithKey (\o _ -> onSightLine here direction o) seats

106 where seatsInDirection = filter (onSightLine here direction) $ M.keys seats

107 sortedSeats = sortOn (manhattan here) seatsInDirection

109 closestInSight :: Seats -> Position -> (S.Set Position)

110 closestInSight seats here = S.fromList

111 $ concatMap (closestInDirection seats here) [d | d <- [Up .. UpLeft]]

113 allSightNeighbourhoods :: Seats -> Neighbourhood

114 allSightNeighbourhoods seats = M.mapWithKey (\h _ -> closestInSight seats h) seats

117 readGrid :: String -> (Seats, Position)

118 readGrid input = (seats, (maxR, maxC))

119 where seats = M.fromList $ concat

120 [ [((r, c), Empty) | (t, c) <- zip row [0..], t == 'L']

121 | (row, r) <- zip rows [0..] ]

122 rows = lines input

123 maxC = (length $ head rows) - 1

124 maxR = (length rows) - 1

126 showGrid seats (maxR, maxC) =

127 unlines $ [ concat [showSeat (r, c) | c <- [0..maxC] ] | r <- [0..maxR]]

128 where showSeat here = show $ M.findWithDefault Floor here seats