run.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeApplications #-}
import AoC
import AoC.Grid

import Control.Monad (guard)
import Data.List (groupBy, sortBy)
import Data.Ord (comparing)

import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Set (Set)
import qualified Data.Set as Set

type N = Int

type Input = (MapGrid Char, (N, N))

parseAll :: String -> Input
parseAll input =
  let g = parseMapGrid id input
  in
    ( g
    , maximum $ HashMap.keys g)

-- TODO: Extract to AoC
pairs :: [a] -> [(a, a)]
pairs l = go l l
  where go [] [] = []
        go (_:xs) [] = go xs xs
        go (x:xs) (y:ys) = (x, y):go (x:xs) ys

antinodes :: [((N, N), a)] -> [V2 N]
antinodes l = do
  ((a, _), (b, _)) <- pairs l
  guard $ a /= b
  let av = v2 a
      bv = v2 b
      d  = av - bv
  [ bv + 2 * d,
    av - 2 * d ]

letters :: MapGrid Char -> [[((N, N), Char)]]
letters g = groupBy (\(_, v1) (_, v2) -> v1 == v2)
          $ sortBy (comparing snd)
          [(k, v) | (k, v) <- HashMap.toList g
                  , v /= '.'
                  ]

part1 :: Input -> Int
part1 (g, bounds) =
  length
  . Set.fromList
  . filter (inBounds bounds)
  . concatMap antinodes
  $ letters g


inBounds :: (N, N) -> V2 N -> Bool
inBounds (w, h) (V2 (x, y)) =
  0 <= x && x <= w && 0 <= y && y <= h

antinodes' :: (N, N) -> [((N, N), a)] -> [V2 N]
antinodes' bounds l = do
  ((a, _), (b, _)) <- pairs l
  guard $ a /= b
  let av = v2 a
      bv = v2 b
      d  = av - bv
  concat [ takeWhile (inBounds bounds) $ map (\k -> bv + fromInteger k * d) [0..]
         , takeWhile (inBounds bounds) $ map (\k -> av - fromInteger k * d) [0..]
         ]

part2 :: Input -> Int
part2 (g, bounds) =
  length
  . Set.fromList
  . concatMap (antinodes' bounds)
  $ letters g

main :: IO ()
main = main' "input.txt"

exampleMain :: IO ()
exampleMain = main' "example.txt"

main' :: FilePath -> IO ()
main' file = do
   input <- parseAll <$> readFile file
   print (part1 input)
   print (part2 input)