61 lines
2 KiB
Haskell
61 lines
2 KiB
Haskell
import Debug.Trace
|
|
import Data.List
|
|
import Data.Maybe
|
|
import Data.Graph
|
|
|
|
traceOut :: Show a => a -> a
|
|
traceOut x = traceShow x x
|
|
|
|
(***) :: (a -> a') -> (b -> b') -> (a, b) -> (a', b')
|
|
(f *** g) (x, y) = (f x, g y)
|
|
|
|
addFirst :: [[a]] -> a -> [[a]]
|
|
addFirst [] x = [[x]]
|
|
addFirst (y:ys) x = (x : y) : ys
|
|
|
|
split :: Eq a => [a] -> [a] -> [[a]]
|
|
split _ [] = []
|
|
split delims (x:xs)
|
|
| x `elem` delims = [] : split delims xs
|
|
| otherwise = addFirst (split delims xs) x
|
|
|
|
enumerate :: [a] -> [(Int, a)]
|
|
enumerate = zip [0..]
|
|
|
|
take3 :: [a] -> Maybe (a, a, a)
|
|
take3 (x:y:z:_) = Just (x, y, z)
|
|
take3 _ = Nothing
|
|
|
|
fst3 :: (a,b,c) -> a
|
|
fst3 (x,_,_) = x
|
|
|
|
parseFile :: String -> Maybe [(Int, Int, Int)]
|
|
parseFile = sequence . map (take3 . map read . split [',']) . lines
|
|
|
|
distanceSq3 :: Num a => (a, a, a) -> (a, a, a) -> a
|
|
distanceSq3 (x, y, z) (p, q, w) = (x-p)^2 + (y-q)^2 + (z-w)^2
|
|
|
|
cartesianProductNeq :: Eq a => [a] -> [a] -> [(a, a)]
|
|
cartesianProductNeq xs ys = [(x, y) | x <- xs, y <- ys, x /= y]
|
|
|
|
solve1 :: [(Int, Int, Int)] -> Integer
|
|
solve1 points = foldl (*) 1 $ take 3 $ sortOn negate $ map (toInteger . length) $ components graph
|
|
where
|
|
joins = take 2000 $ sortOn (uncurry distanceSq3) $ cartesianProductNeq points points
|
|
(graph, vToNode, keyToV) = graphFromEdges $ map (\x -> (x, x, getJoins x)) points
|
|
getJoins x = map fst $ filter ((==x) . snd) joins
|
|
|
|
solve2 :: Int -> [(Int, Int, Int)] -> Int
|
|
solve2 k points = (uncurry (*)) $ (fst3 *** fst3) $ fst $ fromJust $ uncons $ reverse joins
|
|
where
|
|
joins = take (2*k) $ sortOn (uncurry distanceSq3) $ cartesianProductNeq points points
|
|
(graph, vToNode, keyToV) = graphFromEdges $ map (\x -> (x, x, getJoins x)) points
|
|
getJoins x = map fst $ filter ((==x) . snd) joins
|
|
|
|
main = do
|
|
fileinp <- readFile "input.txt"
|
|
let input = fromJust $ parseFile fileinp
|
|
let solved1 = solve1 input
|
|
let solved2 = solve2 7883 input -- found by manual binary search
|
|
print solved1
|
|
print solved2
|