{-# LANGUAGE FlexibleContexts, MultiWayIf, TupleSections, TypeApplications, ViewPatterns #-}
module Day22 (day22a, day22b) where
import Control.Monad.State.Strict (evalState, get, put)
import Data.Array.Unboxed (IArray, Ix, UArray, (!), (//), accumArray, assocs, bounds, elems, listArray, rangeSize)
import Data.Function (on)
import qualified Data.Heap as Heap (MinPolicy, insert, singleton, view)
import qualified Data.Map.Strict as Map (empty, insert, lookup)
import Data.List (scanl')
import Text.Megaparsec (MonadParsec, between, parseMaybe)
import Text.Megaparsec.Char (char, newline, string)
import Text.Megaparsec.Char.Lexer (decimal)
parser :: (Integral a, Integral i, MonadParsec e String m) => m (a, (i, i))
parser = (,) <$> between (string "depth: ") newline decimal
<*> between (string "target: ") newline ((,) <$> decimal <*> (char ',' *> decimal))
erosion, risk :: (Integral a) => a -> a -> a
erosion depth = (`mod` 20183) . (+ depth)
risk depth = (`mod` 3) . erosion depth
makeMaze :: (IArray a e, Enum e, Integral e, Ix e) => e -> (e, e) -> a (e, e) e
makeMaze depth target@(_, targetY) = listArray ((0, 0), target)
(scanl acc [0, 48271..] [1..] >>= take (rangeSize (0, targetY))) // [(target, 0)]
where acc prev x = scanl' ((*) `on` erosion depth) (16807 * x) $ drop 1 prev
growMaze :: (IArray a e, Enum e, Integral e, Ix e) => e -> (e, e) -> a (e, e) e -> a (e, e) e
growMaze depth (sizeX, sizeY) prior = top // bottom
where ((0, 0), (priorX@(max sizeX -> sizeX'), priorY@(max sizeY -> sizeY'))) = bounds prior
top = accumArray (flip const) 0 ((0, 0), (sizeX', sizeY')) $ assocs prior ++ concat right
right = scanl accR [((x, 0), 16807 * x) | x <- [priorX + 1..sizeX']] rightCol
rightCol = [(y, prior ! (priorX, y)) | y <- [1..priorY]]
accR prev (y, v) = zip ((, y) <$> [priorX + 1..]) $ drop 1 $ scanl' (-*-) v $ snd <$> prev
bottom = concat $ drop 1 $ scanl accD bottomRow [priorY + 1..sizeY']
bottomRow = [(p, top ! p) | x <- [0..sizeX'], let p = (x, priorY)]
accD prev y = zip ((, y) <$> [0..]) $ scanl' (-*-) (48271 * y) $ snd <$> drop 1 prev
(-*-) = (*) `on` erosion depth
day22a :: String -> Maybe Int
day22a input = do
(depth, target) <- parseMaybe @() (parser @Int @Int) input
return $ sum $ risk depth <$> elems (makeMaze @UArray depth target)
day22b :: String -> Maybe Int
day22b input = do
(depth, target) <- parseMaybe @() (parser @Int @Int) input
flip evalState (makeMaze @UArray depth target) $ do
let bfs (Heap.view -> Just ((_, (w, k@(p, e))), heap), ws)
| ((0, 0), 1) <- k = return $ Just w
| Just w' <- Map.lookup k ws, w' <= Left w = bfs (heap, ws)
| otherwise = neighbors w p e >>= bfs . commit heap (Map.insert k (Left w) ws)
bfs _ = return Nothing
neighbors w p@(x, y) e = do
maze@(bounds -> ((0, 0), (maxX, maxY))) <- get
maze' <- if x < maxX && y < maxY then return maze else
let maze' = growMaze depth (grow x maxX, grow y maxY) maze
in maze' <$ put maze'
return $ do
x' <- [max x 1 - 1..x + 1]
y' <- if x == x' then [max y 1 - 1..y + 1] else [y]
let (p', r) = ((x', y'), risk depth $ maze' ! p')
if p == p' then [(w + 7, (p', 3 - r - e))] else [(w + 1, (p', e)) | r /= e]
commit heap ws = foldr acc (heap, ws)
where acc a@(w, k@(estimate w -> e)) (heap', ws')
| Just w' <- Map.lookup k ws, w' <= Right e = (heap', ws')
| otherwise = (Heap.insert (e, a) heap', Map.insert k (Right e) ws')
bfs (Heap.singleton @Heap.MinPolicy (0, (0, (target, 1))), Map.empty)
where estimate w ((x, y), e) = w + x + y + if e == 1 then 0 else 7
grow n = head . dropWhile (<= n) . iterate (+ 50)