src/advent20/advent20.hs

   1 {-# LANGUAGE NegativeLiterals #-}
   2 {-# LANGUAGE FlexibleContexts #-}
   3 {-# LANGUAGE OverloadedStrings #-}
   4 {-# LANGUAGE TypeFamilies #-}
   5 {-# LANGUAGE BangPatterns #-}
   6
   7 import Data.Text (Text)
   8 import qualified Data.Text as T
   9 import qualified Data.Text.IO as TIO
  10
  11 import Text.Megaparsec hiding (State)
  12 import qualified Text.Megaparsec.Lexer as L
  13 import Text.Megaparsec.Text (Parser)
  14 import qualified Control.Applicative as CA
  15
  16 import qualified Data.Vector as V
  17
  18 import Data.List
  19
  20 import qualified Data.Set as S
  21
  22
  23 type Vec = V.Vector Integer
  24
  25 data Particle = Particle
  26                     { position :: Vec
  27                     , velocity :: Vec
  28                     , acceleration :: Vec
  29                     } deriving (Show, Eq)
  30
  31
  32 main :: IO ()
  33 main = do
  34         text <- TIO.readFile "data/advent20.txt"
  35         let particles = successfulParse text
  36         print $ part1 particles
  37         print $ part2 50 particles
  38
  39
  40 part1 :: [Particle] -> Int
  41 part1 particles = head $ withMinX $ simulate particles
  42
  43 part2 :: Integer -> [Particle] -> Int
  44 part2 n particles = length $ simulateC n particles
  45
  46 simulate :: [Particle] -> [Particle]
  47 simulate particles =
  48     if all quiescent particles && length withMinXs == 1
  49     then particles
  50     else simulate (map step particles)
  51     where withMinXs = withMinX particles
  52
  53
  54 simulateC :: Integer -> [Particle] -> [Particle]
  55 simulateC 0 particles = particles
  56 simulateC t particles = simulateC (t - 1) (map step particles')
  57     where particles' = removeColliders particles
  58
  59
  60 step :: Particle -> Particle
  61 step particle = particle {position = p', velocity = v'}
  62     where pv' = V.zipWith3 updatePV (position particle)
  63                                     (velocity particle)
  64                                     (acceleration particle)
  65           (p', v') = V.unzip pv'
  66           updatePV p v a = (p + v + a, v + a)
  67
  68
  69 -- Checks whether a particle could ever get closer to the origin than it is now.
  70 quiescent :: Particle -> Bool
  71 quiescent particle = and qDimensions
  72     where qDimensions = V.zipWith3 sameSigns (position particle)
  73                                              (velocity particle)
  74                                              (acceleration particle)
  75           sameSigns p v a = if a == 0 && v == 0
  76                             then True
  77                             else if a == 0
  78                                  then signum p == signum v
  79                                  else signum p == signum v
  80                                         && signum v == signum a
  81
  82
  83 withMinX particles = minX `elemIndices` absXs
  84     where absXs = map pAbsX particles
  85           minX = minimum absXs
  86
  87 pAbsX :: Particle -> Integer
  88 pAbsX particle = V.foldl1' (+) $ V.map abs (position particle)
  89
  90
  91 removeColliders particles = particles'
  92     where positions = map position particles
  93           duplicatePositions = S.fromList $ concat
  94                                           $ filter (\g -> length g > 1)
  95                                           $ group
  96                                           $ sort positions
  97           particles' = filter (\p -> not (S.member (position p)
  98                                                    duplicatePositions))
  99                               particles
 100
 101
 102
 103 sc :: Parser ()
 104 sc = L.space (skipSome spaceChar) CA.empty CA.empty
 105
 106 lexeme  = L.lexeme sc
 107
 108 integer       = lexeme L.integer
 109 signedInteger = L.signed sc integer
 110
 111 symbol = L.symbol sc
 112 separator = symbol ", "
 113 comma = symbol ","
 114
 115 particlesP = particleP `sepBy` space
 116 particleP = particlify <$> (symbol "p=" *> vecP <* separator)
 117                        <*> (symbol "v=" *> vecP <* separator)
 118                        <*> (symbol "a=" *> vecP)
 119     where particlify p v a = Particle { position = p
 120                                       , velocity = v
 121                                       , acceleration = a
 122                                       }
 123
 124
 125 vecP = V.fromList <$> between (symbol "<") (symbol ">")
 126                               (signedInteger `sepBy` comma)
 127
 128
 129 successfulParse :: Text -> [Particle]
 130 successfulParse input =
 131         case parse particlesP "input" input of
 132                 Left  _error -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 133                 Right instructions  -> instructions