git.haldean.org aoc-2018 / master 06 / Aoc06.hs
master

Tree @master (Download .tar.gz)

Aoc06.hs @masterraw · history · blame

import Control.Arrow
import Data.Function
import qualified Data.List as L
import qualified Data.Set as S
import Text.Parsec
import Text.Parsec.String

type Coord = (Int, Int)
type Base = (Int, Int, Int) -- x, y, id
type BasedCoord = (Int, Int, Int) -- x, y, id

parseInt = fmap (read :: String -> Int) $ many1 digit
parseBaseLoc :: Parser (Int, Int)
parseBaseLoc = do
  x <- parseInt
  _ <- string ", "
  y <- parseInt
  return (x, y)

readBases :: String -> Either ParseError [Base]
readBases s = case parse (sepEndBy parseBaseLoc (char '\n')) "input.txt" s of
  Right coords -> Right $ zipWith (\(x, y) c -> (x, y, c)) coords [1..]
  Left err -> Left err

gridSize = 400
allCoords :: [Coord]
allCoords =
  let xs = enumFromTo 0 gridSize
      ys = enumFromTo 0 gridSize
  in concat $ map (\xx -> zip (repeat xx) ys) xs

minBy :: Ord a => Ord b => (a -> b) -> [a] -> a
minBy f as = snd $ L.minimum $ zip (map f as) as

manhattan (x1, y1) (x2, y2) = abs (x1 - x2) + abs (y1 - y2)

mapToNearest :: [Base] -> Coord -> BasedCoord
mapToNearest bs c@(x, y) = let (_, _, b) = minBy (\(xb, yb, b) -> manhattan c (xb, yb)) bs in (x, y, b)

baseId :: Base -> Int
baseId (_, _, b) = b
baseCoord :: Base -> Coord
baseCoord (x, y, _) = (x, y)

findInfBases :: [BasedCoord] -> S.Set Int
findInfBases bcs =
  filter (\(x, y, _) -> x == 0 || x == gridSize || y == 0 || y == gridSize) bcs
  & map baseId
  & S.fromList

findAreas :: [Base] -> [(Int, Base)]
findAreas bs =
  let bcs = map (mapToNearest bs) allCoords
  in let infBases = findInfBases bcs
  in let goodBases = filter (\(_, _, b) -> not $ S.member b infBases) bs
  in map (\base@(_, _, b) -> (length $ filter (\(_, _, bcb) -> b == bcb) bcs, base)) goodBases

findSumDist :: [Base] -> Coord -> Int
findSumDist bs c = sum $ map (manhattan c . baseCoord) bs

countGoodPoints :: [Base] -> Int
countGoodPoints bs = length $ filter (\c -> findSumDist bs c < 10000) allCoords

main = readFile "input.txt" >>= \str -> case readBases str of
  Right bases -> do
    putStrLn $ (++) "1: " $ show $ L.maximum $ findAreas bases
    putStrLn $ (++) "2: " $ show $ countGoodPoints bases
  Left err -> print err