src/advent03/advent03.hs

   1 import Data.List (tails)
   2 import qualified Data.HashMap.Strict as M
   3 import Data.HashMap.Strict ((!))
   4 import Control.Monad.State.Lazy
   5
   6 type Location = (Int, Int)
   7 type Memory = M.HashMap Location Int
   8
   9 target :: Int
  10 target = 347991
  11
  12 main :: IO ()
  13 main = do
  14         print $ part1
  15         print $ part2
  16
  17
  18 diagonal :: Int -> [Int]
  19 diagonal n = scanl (+) 1 $ scanl (+) n $ repeat 8
  20 dr = diagonal 8
  21 ul = diagonal 4
  22 ur = diagonal 2
  23 dl = diagonal 6
  24
  25
  26 interleave :: [[a]] -> [a]
  27 interleave ([]:_) = []
  28 interleave xss = map head xss ++ interleave (map tail xss)
  29
  30
  31 countedDiags = interleave [(zip [0..] ur), (zip [0..] ul), (zip [0..] dl), (zip [0..] dr)]
  32
  33 part1 = let corners =  head $ dropWhile ((< target) . snd . head . tail) $ tails countedDiags
  34             (pcd, pcv) = head corners
  35             (ncd, ncv) = head $ tail corners
  36             result = if pcd == ncd
  37                         then if (target - pcv + 2) < ncv - target
  38                              then pcd * 2 - (target - pcv)
  39                              else ncd * 2 - (ncv - target)
  40                      else if (target - pcv + 1) < ncv - target
  41                              then pcd * 2 - (target - pcv) + 2
  42                              else ncd * 2 - (ncv - target)
  43     in result
  44
  45
  46 part2 = (!) memoryValues $ head $ dropWhile (\l -> memoryValues!l <= target) locations
  47     where memoryValues = execState (updateMemory (take 100 $ drop 1 locations)) emptyMemory
  48
  49 up    (a, b) = (a, b + 1)
  50 down  (a, b) = (a, b - 1)
  51 left  (a, b) = (a - 1, b)
  52 right (a, b) = (a + 1, b)
  53 directions = [right, up, left, down]
  54
  55 locations = scanl (\c f -> f c) (0,0) $ concat $ zipWith replicate steps (cycle directions)
  56     where
  57         steps = concat $ zipWith (\a b -> [a,b]) [1..] [1..]
  58
  59 adjacentMap (r, c) = M.filterWithKey adjacent
  60     where adjacent k _ = abs (fst k - r) <= 1 && abs (snd k - c) <= 1
  61
  62 adjacentMapSum here = M.foldr (+) 0 . adjacentMap here
  63
  64
  65 emptyMemory = M.singleton (0, 0) 1
  66
  67 updateMemoryOnce :: Location -> State Memory Int
  68 updateMemoryOnce here =
  69     do m0 <- get
  70        let total = adjacentMapSum here m0
  71        put (M.insert here total m0)
  72        return total
  73
  74 updateMemory :: [Location] -> State Memory Int
  75 updateMemory [] = do return 0
  76 updateMemory (l:ls) =
  77     do  updateMemoryOnce l
  78         updateMemory ls
  79