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 $ S.size trees
  13       -- print trees
  14       -- print $ maxCorner
  15       -- print $ take 25 $ visitedPlaces (1, 3) maxCorner
  16       -- print $ takeWhile (withinTrees maxCorner) $ visitedPlaces (1, 3) maxCorner
  17       print $ part1 trees maxCorner
  18       print $ part2 trees maxCorner
  19
  20
  21 readGrid :: String -> (Trees, Position)
  22 readGrid input = ( S.fromList [ (r, c)
  23                               | r <- [0..maxR], c <- [0..maxC]
  24                               , (grid!!r)!!c == '#']
  25                  , (maxR, maxC)
  26                  )
  27     where grid = lines input
  28           maxC = (length $ head grid) - 1
  29           maxR = (length grid) - 1
  30
  31
  32 part1 trees maxCorner = countEncounteredTrees trees maxCorner (1, 3)
  33
  34 part2 trees maxCorner = foldr1 (*) $ map cet [(1, 1), (1, 3), (1, 5), (1, 7), (2, 1)]
  35   where cet = countEncounteredTrees trees maxCorner
  36
  37 countEncounteredTrees trees maxCorner delta = S.size $ S.intersection trees visited
  38   where visited = S.fromList $ takeWhile (withinTrees maxCorner) $ visitedPlaces delta maxCorner
  39
  40
  41 visitedPlaces :: Position -> Position -> [Position]
  42 visitedPlaces (dr, dc) (_maxR, maxC) = iterate wrappingStep (0, 0)
  43   where wrappingStep (r, c) = (r + dr, (c + dc) `mod` (maxC + 1))
  44
  45 withinTrees (maxR, maxC) (r, c) = r >= 0 && r <= maxR && c >= 0 && c <= maxC
  46