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