ac61359be7a29617c50f1377839d67a5c267553a
1 import Data.Text (Text)
2 import qualified Data.Text.IO as TIO
4 import Data.Void (Void)
6 import Text.Megaparsec hiding (State)
7 import Text.Megaparsec.Char
8 import qualified Text.Megaparsec.Char.Lexer as L
9 import qualified Control.Applicative as CA
11 import Linear (V3(..), V1(..), (^+^), (^-^))
13 import qualified Data.Vector as V
15 class (Eq a) => NVec a where
16 (^+^^) :: a -> a -> a
17 (^-^^) :: a -> a -> a
18 nvZero :: a
19 nvSignum :: a -> a
20 nvAbsSum :: a -> Integer
21 instance NVec (V1 Integer) where
22 x ^+^^ y = x ^+^ y
23 x ^-^^ y = x ^-^ y
24 nvZero = V1 0
25 nvSignum (V1 x) = V1 (signum x)
26 nvAbsSum (V1 x) = abs x
27 instance NVec (V3 Integer) where
28 x ^+^^ y = x ^+^ y
29 x ^-^^ y = x ^-^ y
30 nvZero = V3 0 0 0
31 nvSignum (V3 x y z) = V3 (signum x) (signum y) (signum z)
32 nvAbsSum (V3 x y z) = (abs x) + (abs y) + (abs z)
35 data Planet a = Planet { _pos :: a , _vel :: a} deriving (Show, Eq, Ord)
36 type Planets a = V.Vector (Planet a)
39 main :: IO ()
40 main = do
42 let planetsT = successfulParse text
43 let planets = enplanet planetsT
44 -- print planets
45 print \$ part1 planets
46 print \$ part2 planets
48 part1 :: Planets (V3 Integer) -> Integer
49 part1 = systemEnergy . head . drop 1000 . simulate
51 part2 :: Planets (V3 Integer) -> Integer
52 part2 planets = period
53 where planetDimensions = unzipPlanets planets
54 simCounts = map countSimulate planetDimensions
55 period = foldl lcm 1 simCounts
58 enplanet :: (NVec a) => [a] -> Planets a
59 enplanet = V.fromList . map mkPlanet
61 mkPlanet :: (NVec a) => a -> Planet a
62 mkPlanet p = Planet {_pos = p, _vel = nvZero}
65 unzipPlanets :: V.Vector (Planet (V3 Integer)) -> [V.Vector (Planet (V1 Integer))]
66 unzipPlanets = dimensionSlice . V.map unzipPlanet
68 unzipPlanet :: Planet (V3 Integer) -> [Planet (V1 Integer)]
69 unzipPlanet planet = map mkPlanet posVecs
70 where posVecs = unzipVec \$ _pos planet
73 unzipVec :: V3 Integer -> [V1 Integer]
74 unzipVec (V3 x y z) = [V1 x, V1 y, V1 z]
76 dimensionSlice :: (NVec a) => V.Vector [Planet a] -> [Planets a]
77 dimensionSlice slicedPlanets = [sliceDim d | d <- [0..2]]
78 where sliceDim d = V.map (!!d) slicedPlanets
82 simulate :: (NVec a) => Planets a -> [Planets a]
83 simulate = iterate simulationStep
85 countSimulate :: (NVec a) => Planets a -> Integer
86 countSimulate planets0 = go (simulationStep planets0) 1
87 where go planets n
88 | planets0 == planets = n
89 | otherwise = go (simulationStep planets) (n + 1)
92 simulationStep :: (NVec a) => Planets a -> Planets a
93 simulationStep planets = planets''
94 where planets' = applyGravity planets
95 planets'' = applyVelocity planets'
98 gravity :: (NVec a) => a -> a
99 gravity v = nvSignum v
101 applyGravity :: (NVec a) => Planets a -> Planets a
102 applyGravity planets = V.map (applyGravityHere planets) planets
104 applyGravityHere :: (NVec a) => Planets a -> Planet a -> Planet a
105 applyGravityHere planets here = V.foldl' updateGravity here planets
107 updateGravity :: (NVec a) => Planet a -> Planet a -> Planet a
108 updateGravity here there = here { _vel = vel'}
109 where vel = _vel here
110 vel' = vel ^+^^ gravity ((_pos there) ^-^^ (_pos here))
112 applyVelocity :: (NVec a) => Planets a -> Planets a
113 applyVelocity = V.map applyVelocityHere
115 applyVelocityHere :: (NVec a) => Planet a -> Planet a
116 applyVelocityHere here = here {_pos = (_pos here) ^+^^ (_vel here)}
121 -- absSum (Vec1 (V1 x)) = (abs x)
122 -- absSum (Vec3 (V3 x y z)) = (abs x) + (abs y) + (abs z)
124 potentalEnergy planet = nvAbsSum \$ _pos planet
125 kineticEnergy planet = nvAbsSum \$ _vel planet
126 totalEnergy planet = (potentalEnergy planet) * (kineticEnergy planet)
128 systemEnergy = (V.foldl' (+) 0) . (V.map totalEnergy)
132 -- Parse the input file
133 type Parser = Parsec Void Text
135 sc :: Parser ()
136 sc = L.space (skipSome spaceChar) CA.empty CA.empty
137 -- sc = L.space (skipSome (char ' ')) CA.empty CA.empty
139 lexeme = L.lexeme sc
140 integer = lexeme L.decimal
141 signedInteger = L.signed sc integer
142 symb = L.symbol sc
143 equalP = symb "="
144 commaP = symb ","
145 identifierP = some alphaNumChar <* sc
146 openBracketP = symb "<"
147 closeBracketP = symb ">"
149 planetsP = many planetP
151 planetP = (between openBracketP closeBracketP) coordsP
153 coordsP = envector <\$> (coordP `sepBy` commaP)
154 where envector [x, y, z] = V3 x y z
155 coordP = identifierP *> equalP *> signedInteger
158 -- successfulParse :: Text -> [Vec]
159 successfulParse :: Text -> [V3 Integer]
160 successfulParse input =
161 case parse planetsP "input" input of
162 Left _err -> [] -- TIO.putStr \$ T.pack \$ parseErrorPretty err
163 Right planets -> planets