{-# Language QuasiQuotes, ImportQualifiedPost #-}
module Main where
import Advent (format, minimumMaybe)
import Data.Foldable (toList)
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.List (foldl', tails, intersect)
import Data.Map qualified as Map
import Linear (V2(V2), V3(V3), quadrance)
import Linear.Matrix as LM (transpose)
main :: IO ()
IO ()
main =
do [(Int, Int, Int, Int, Int, Int, Int, Int, Int)]
input <- [format|2017 20 (p=<%d,%d,%d>, v=<%d,%d,%d>, a=<%d,%d,%d>%n)*|]
let particles :: [V3 (V3 Double)]
particles = [(Int -> Double) -> V3 Int -> V3 Double
forall a b. (a -> b) -> V3 a -> V3 b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (V3 Int -> V3 Double) -> V3 (V3 Int) -> V3 (V3 Double)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> V3 Int -> V3 Int -> V3 Int -> V3 (V3 Int)
forall a. a -> a -> a -> V3 a
V3 (Int -> Int -> Int -> V3 Int
forall a. a -> a -> a -> V3 a
V3 Int
p1 Int
p2 Int
p3) (Int -> Int -> Int -> V3 Int
forall a. a -> a -> a -> V3 a
V3 Int
v1 Int
v2 Int
v3) (Int -> Int -> Int -> V3 Int
forall a. a -> a -> a -> V3 a
V3 Int
a1 Int
a2 Int
a3) | (Int
p1,Int
p2,Int
p3,Int
v1,Int
v2,Int
v3,Int
a1,Int
a2,Int
a3) <- [(Int, Int, Int, Int, Int, Int, Int, Int, Int)]
input]
Int -> IO ()
forall a. Show a => a -> IO ()
print ([V3 (V3 Double)] -> Int
part1 [V3 (V3 Double)]
particles)
Int -> IO ()
forall a. Show a => a -> IO ()
print ([V3 (V3 Double)] -> Int
part2 [V3 (V3 Double)]
particles)
type Particle = V3 (V3 Double)
part1 :: [Particle] -> Int
part1 :: [V3 (V3 Double)] -> Int
part1 = (V3 (V3 Double) -> V3 Double) -> [V3 (V3 Double)] -> Int
forall b a. Ord b => (a -> b) -> [a] -> Int
minimumIndexOn ((V3 Double -> Double) -> V3 (V3 Double) -> V3 Double
forall a b. (a -> b) -> V3 a -> V3 b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap V3 Double -> Double
forall a. Num a => V3 a -> a
forall (f :: * -> *) a. (Metric f, Num a) => f a -> a
quadrance)
part2 :: [Particle] -> Int
part2 :: [V3 (V3 Double)] -> Int
part2 [V3 (V3 Double)]
ps = [V3 (V3 Double)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [V3 (V3 Double)]
ps Int -> Int -> Int
forall a. Num a => a -> a -> a
- IntSet -> Int
IntSet.size ([[V2 Int]] -> IntSet
process ([V3 (V3 Double)] -> [[V2 Int]]
toEvents [V3 (V3 Double)]
ps))
stepParticle :: Particle -> Particle
stepParticle :: V3 (V3 Double) -> V3 (V3 Double)
stepParticle (V3 V3 Double
a V3 Double
v V3 Double
s) = V3 Double -> V3 Double -> V3 Double -> V3 (V3 Double)
forall a. a -> a -> a -> V3 a
V3 V3 Double
a V3 Double
v' V3 Double
s'
where
v' :: V3 Double
v' = V3 Double
a V3 Double -> V3 Double -> V3 Double
forall a. Num a => a -> a -> a
+ V3 Double
v
s' :: V3 Double
s' = V3 Double
s V3 Double -> V3 Double -> V3 Double
forall a. Num a => a -> a -> a
+ V3 Double
v'
minimumIndexOn :: Ord b => (a -> b) -> [a] -> Int
minimumIndexOn :: forall b a. Ord b => (a -> b) -> [a] -> Int
minimumIndexOn a -> b
f [a]
xs = (b, Int) -> Int
forall a b. (a, b) -> b
snd ([(b, Int)] -> (b, Int)
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum ((a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f [a]
xs [b] -> [Int] -> [(b, Int)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Int
0..]))
collide :: V3 (V3 Double) -> V3 (V3 Double) -> Maybe Double
collide :: V3 (V3 Double) -> V3 (V3 Double) -> Maybe Double
collide V3 (V3 Double)
p1 V3 (V3 Double)
p2
= [Double] -> Maybe Double
forall a. Ord a => [a] -> Maybe a
minimumMaybe
([Double] -> Maybe Double) -> [Double] -> Maybe Double
forall a b. (a -> b) -> a -> b
$ (Double -> Bool) -> [Double] -> [Double]
forall a. (a -> Bool) -> [a] -> [a]
filter (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0)
([Double] -> [Double]) -> [Double] -> [Double]
forall a b. (a -> b) -> a -> b
$ ([Double] -> [Double] -> [Double]) -> V3 [Double] -> [Double]
forall a. (a -> a -> a) -> V3 a -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 [Double] -> [Double] -> [Double]
forall a. Eq a => [a] -> [a] -> [a]
intersect
(V3 [Double] -> [Double]) -> V3 [Double] -> [Double]
forall a b. (a -> b) -> a -> b
$ (V3 Double -> [Double]) -> V3 (V3 Double) -> V3 [Double]
forall a b. (a -> b) -> V3 a -> V3 b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (V3 Double -> [Double]
zeros (V3 Double -> [Double])
-> (V3 Double -> V3 Double) -> V3 Double -> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. V3 Double -> V3 Double
toPoly)
(V3 (V3 Double) -> V3 [Double]) -> V3 (V3 Double) -> V3 [Double]
forall a b. (a -> b) -> a -> b
$ V3 (V3 Double)
p1 V3 (V3 Double) -> V3 (V3 Double) -> V3 (V3 Double)
forall a. Num a => a -> a -> a
- V3 (V3 Double)
p2
toPoly :: V3 Double -> V3 Double
toPoly :: V3 Double -> V3 Double
toPoly (V3 Double
a Double
v Double
s) = Double -> Double -> Double -> V3 Double
forall a. a -> a -> a -> V3 a
V3 (Double
aDouble -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
2) (Double
aDouble -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
v) Double
s
zeros ::
V3 Double ->
[Double]
zeros :: V3 Double -> [Double]
zeros (V3 Double
0 Double
0 Double
0) = [Double
0]
zeros (V3 Double
0 Double
0 Double
_) = []
zeros (V3 Double
0 Double
b Double
c) = [-Double
cDouble -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
b]
zeros (V3 Double
a Double
b Double
c) =
case Double -> Double -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Double
z Double
0 of
Ordering
LT -> []
Ordering
EQ -> [-Double
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
a]
Ordering
GT -> [(-Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
sqrt Double
z) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
a), (-Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
sqrt Double
z) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
a)]
where
z :: Double
z = Double
bDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
4Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
aDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
c
toEvents :: [Particle] -> [[V2 Int]]
toEvents :: [V3 (V3 Double)] -> [[V2 Int]]
toEvents [V3 (V3 Double)]
ps
= Map Double [V2 Int] -> [[V2 Int]]
forall a. Map Double a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList
(Map Double [V2 Int] -> [[V2 Int]])
-> Map Double [V2 Int] -> [[V2 Int]]
forall a b. (a -> b) -> a -> b
$ ([V2 Int] -> [V2 Int] -> [V2 Int])
-> [(Double, [V2 Int])] -> Map Double [V2 Int]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [V2 Int] -> [V2 Int] -> [V2 Int]
forall a. [a] -> [a] -> [a]
(++)
[ (Double
t,[Int -> Int -> V2 Int
forall a. a -> a -> V2 a
V2 Int
i Int
j])
| (Int
i,V3 (V3 Double)
v1):[(Int, V3 (V3 Double))]
vs <- [(Int, V3 (V3 Double))] -> [[(Int, V3 (V3 Double))]]
forall a. [a] -> [[a]]
tails ([Int
0..] [Int] -> [V3 (V3 Double)] -> [(Int, V3 (V3 Double))]
forall a b. [a] -> [b] -> [(a, b)]
`zip` (V3 (V3 Double) -> V3 (V3 Double))
-> [V3 (V3 Double)] -> [V3 (V3 Double)]
forall a b. (a -> b) -> [a] -> [b]
map V3 (V3 Double) -> V3 (V3 Double)
forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
LM.transpose [V3 (V3 Double)]
ps)
, (Int
j,V3 (V3 Double)
v2) <- [(Int, V3 (V3 Double))]
vs
, Just Double
t <- [V3 (V3 Double) -> V3 (V3 Double) -> Maybe Double
collide V3 (V3 Double)
v1 V3 (V3 Double)
v2] ]
process :: [[V2 Int]] -> IntSet
process :: [[V2 Int]] -> IntSet
process = (IntSet -> [V2 Int] -> IntSet) -> IntSet -> [[V2 Int]] -> IntSet
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IntSet -> [V2 Int] -> IntSet
runGeneration IntSet
IntSet.empty
runGeneration ::
IntSet ->
[V2 Int] ->
IntSet
runGeneration :: IntSet -> [V2 Int] -> IntSet
runGeneration IntSet
dead [V2 Int]
xs =
(IntSet -> Int -> IntSet) -> IntSet -> [Int] -> IntSet
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Int -> IntSet -> IntSet) -> IntSet -> Int -> IntSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> IntSet -> IntSet
IntSet.insert) IntSet
dead
[ Int
c | V2 Int
x <- [V2 Int]
xs
, (Int -> Bool) -> V2 Int -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> IntSet -> Bool
`IntSet.notMember` IntSet
dead) V2 Int
x
, Int
c <- V2 Int -> [Int]
forall a. V2 a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList V2 Int
x ]