advent03/src/advent03.hs

   1 -- import Debug.Trace
   2
   3 import qualified Data.Set as S
   4
   5 type Position = (Int, Int)
   6 type Trees = S.Set Position
   7
   8 main :: IO ()
   9 main =
  10   do  text <- readFile "data/advent03.txt"
  11       let (trees, maxCorner) = readGrid text
  12       print $ part1 trees maxCorner
  13       print $ part2 trees maxCorner
  14
  15
  16 readGrid :: String -> (Trees, Position)
  17 readGrid input = (trees, (maxR, maxC))
  18   where trees = S.fromList $ concat
  19                   [ [(r, c) | (t, c) <- zip row [0..], t == '#']
  20                   | (row, r) <- zip rows [0..] ]
  21         rows = lines input
  22         maxC = (length $ head rows) - 1
  23         maxR = (length rows) - 1
  24
  25 part1 trees maxCorner = countEncounteredTrees trees maxCorner (1, 3)
  26
  27 part2 trees maxCorner = product $ map cet [(1, 1), (1, 3), (1, 5), (1, 7), (2, 1)]
  28   where cet = countEncounteredTrees trees maxCorner
  29
  30 countEncounteredTrees trees maxCorner delta = S.size $ S.intersection trees visited
  31   where visited = S.fromList $ takeWhile (withinTrees maxCorner) $ visitedPlaces delta maxCorner
  32
  33 visitedPlaces :: Position -> Position -> [Position]
  34 visitedPlaces (dr, dc) (_maxR, maxC) = iterate wrappingStep (0, 0)
  35   where wrappingStep (r, c) = (r + dr, (c + dc) `mod` (maxC + 1))
  36
  37 withinTrees (maxR, maxC) (r, c) = r >= 0 && r <= maxR && c >= 0 && c <= maxC