advent19/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2021/12/21/advent-of-code-2021-day-19/
   2
   3 import Data.Text ()
   4 import qualified Data.Text.IO as TIO
   5 import Data.Attoparsec.Text hiding (take, takeWhile)
   6
   7 import Linear (V3(..), (^+^), (^-^))
   8 import qualified Data.Set as S
   9 import qualified Data.MultiSet as MS
  10 import Data.Monoid
  11 import Data.Maybe
  12 import Data.List
  13 import Control.Monad
  14
  15 type Coord = V3 Int
  16
  17 data Scanner = Scanner
  18   { scannerName :: Int
  19   , beacons :: [Coord]
  20   , transformation :: Endo Coord
  21   , signature :: MS.MultiSet Int
  22   }
  23   deriving (Show)
  24
  25 instance Eq Scanner where
  26   s1 == s2 = (scannerName s1) == (scannerName s2)
  27
  28 data Reconstruction = Reconstruction
  29   { found :: [Scanner] -- these have had the transform applied to the beacons
  30   , working :: [Scanner] -- these have had the transform applied to the beacons
  31   , waiting :: [Scanner] -- these are as discovered
  32   }
  33   deriving (Show)
  34
  35
  36 -- Transformations
  37
  38 type Transform = Endo Coord
  39
  40 instance Show Transform where
  41   -- show c = show $ appEndo c (V3 1 2 3)
  42   show c = show $ appEndo c (V3 0 0 0)
  43
  44 nullTrans = Endo id
  45 rotX = Endo \(V3 x y z) -> V3    x (- z)   y
  46 rotY = Endo \(V3 x y z) -> V3    z    y (- x)
  47 rotZ = Endo \(V3 x y z) -> V3 (- y)   x    z
  48 translate v = Endo (v ^+^)
  49
  50 rotations :: [Transform]
  51 rotations = [a <> b | a <- ras, b <- rbs]
  52   where ras = [ nullTrans, rotY, rotY <> rotY, rotY <> rotY <> rotY
  53               , rotZ, rotZ <> rotZ <> rotZ]
  54         rbs = [nullTrans, rotX, rotX <> rotX, rotX <> rotX <> rotX]
  55
  56 -- Main
  57
  58 main :: IO ()
  59 main =
  60   do  text <- TIO.readFile "data/advent19.txt"
  61       let scanners = successfulParse text
  62       let rec0 = mkReconstruction scanners
  63       let rec = reconstruct rec0
  64       let transScanners = found rec
  65       print $ part1 transScanners
  66       print $ part2 transScanners
  67
  68 part1 :: [Scanner] -> Int
  69 part1 scanners = S.size $ S.unions $ map (S.fromList . beacons) scanners
  70
  71 part2 :: [Scanner] -> Int
  72 part2 scanners = maximum [manhattan (a ^-^ b) | a <- origins, b <- origins]
  73   where extractOrigin sc = appEndo (transformation sc) (V3 0 0 0)
  74         origins = map extractOrigin scanners
  75         manhattan (V3 x y z) = (abs x) + (abs y) + (abs z)
  76
  77 sign :: [Coord] -> MS.MultiSet Int
  78 sign bcns = MS.fromList [pythag (a ^-^ b) | a <- bcns, b <- bcns, a < b]
  79   where pythag (V3 x y z) = x^2 + y^2 + z^2
  80
  81 vagueMatch :: Scanner -> Scanner -> Bool
  82 vagueMatch scanner1 scanner2 = s >= (12 * 11) `div` 2
  83   where s = MS.size $ MS.intersection (signature scanner1) (signature scanner2)
  84
  85 matchingTransform :: Scanner -> Scanner -> Maybe Transform
  86 matchingTransform s1 s2 = listToMaybe $ matchingTransformAll s1 s2
  87
  88 matchingTransformAll :: Scanner -> Scanner -> [Transform]
  89 matchingTransformAll scanner1 scanner2 =
  90   do  let beacons1 = beacons scanner1
  91       let beacons2 = beacons scanner2
  92       rot <- rotations
  93       b1 <- beacons1
  94       b2 <- beacons2
  95       let t = b1 ^-^ (appEndo rot b2)
  96       let translation = translate t
  97       let transB2 = map (appEndo (translation <> rot)) beacons2
  98       guard $ (length $ intersect beacons1 transB2) >= 12
  99       return (translation <> rot)
 100
 101
 102 mkReconstruction :: [Scanner] -> Reconstruction
 103 mkReconstruction (s:ss) = Reconstruction {found = [], working = [s], waiting = ss}
 104
 105 reconstruct :: Reconstruction -> Reconstruction
 106 reconstruct r
 107   -- | waiting r == [] = Reconstruction { found = (found r) ++ (working r), working = [], waiting = []}
 108   | working r == [] = r
 109   | otherwise = reconstruct $ reconstructStep r
 110
 111 reconstructStep :: Reconstruction -> Reconstruction
 112 reconstructStep Reconstruction{..} =
 113   Reconstruction { found = current : found
 114                  , working = workers ++ newWorkers
 115                  , waiting = waiting'
 116                  }
 117   where (current:workers) = working
 118         possMatches = filter (vagueMatch current) waiting
 119         matches = filter (isJust . snd) $ zip possMatches $ map (matchingTransform current) possMatches
 120         waiting' = waiting \\ (map fst matches)
 121         newWorkers = map (transformScanner) matches
 122
 123 transformScanner :: (Scanner, Maybe Transform) -> Scanner
 124 transformScanner (Scanner{..}, trans) =
 125   Scanner { beacons = map (appEndo $ fromJust trans) beacons
 126           , transformation = fromJust trans
 127           , ..}
 128
 129
 130 -- Parse the input file
 131
 132 scannersP = scannerP `sepBy` blankLines
 133 scannerP = scannerify <$> nameP <*> beaconsP
 134   where
 135     scannerify name beacons =
 136       Scanner { scannerName = name
 137               , beacons = beacons
 138               , transformation = nullTrans
 139               , signature = sign beacons
 140               }
 141
 142 nameP = ("--- scanner " *>) decimal <* " ---" <* endOfLine
 143
 144 beaconsP = beaconP `sepBy` endOfLine
 145 beaconP = V3 <$> (signed decimal <* ",") <*> (signed decimal <* ",") <*> (signed decimal)
 146
 147 blankLines = many1 endOfLine
 148
 149 -- successfulParse :: Text -> (Integer, [Maybe Integer])
 150 successfulParse input =
 151   case parseOnly scannersP input of
 152     Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
 153     Right scanners -> scanners