08.hs
1import Data.List (find, sortBy)
2import Data.List.Split (splitOn)
3import Data.Ord (comparing)
4import qualified Data.Set as S
5
6parse = map (map read . splitOn ",") . lines
7
8dist [x, y, z] [x', y', z'] = sqrt $ (x - x') ^ 2 + (y - y') ^ 2 + (z - z') ^ 2
9
10pairwise ls = sortBy (comparing snd) [((i, i'), dist p p') | (i, p) <- zip [0 ..] ls, (i', p') <- zip [0 ..] ls, i' > i]
11
12circuits ls limit = foldl insert [] ordered
13 where
14 ordered = map fst . take limit $ pairwise ls
15 insert (c : cs) (i, i') = if any (`S.member` c) [i, i'] then mergeSets ((c `S.union` S.fromList [i, i']) : cs) else c : insert cs (i, i')
16 insert [] (i, i') = [S.fromList [i, i']]
17
18circuits2 ls = (ls !! p, ls !! p')
19 where
20 Just ((p, p'), _) = find (\(pt, res) -> length ls == maximum (map S.size res)) $ tail $ scanl insert ((0, 0), []) ordered
21 ordered = map fst $ pairwise ls
22 insert :: ((Int, Int), [S.Set Int]) -> (Int, Int) -> ((Int, Int), [S.Set Int])
23 insert (_, (c : cs)) (i, i') =
24 if any (`S.member` c) [i, i']
25 then
26 let merged = mergeSets ((c `S.union` S.fromList [i, i']) : cs)
27 pt = if (maximum $ map S.size merged) == length ls then (i, i') else (0, 0)
28 in (pt, merged)
29 else ((0, 0), c : snd (insert ((0, 0), cs) (i, i')))
30 insert (_, []) (i, i') = ((0, 0), [S.fromList [i, i']])
31
32mergeSets [] = []
33mergeSets (s : ss) = mergeSets' [s] ss
34 where
35 mergeSets' acc [] = acc
36 mergeSets' acc (x : xs) =
37 case findIntersecting x acc of
38 Nothing -> mergeSets' (x : acc) xs
39 Just (intersecting, rest) ->
40 mergeSets' (S.unions (x : intersecting) : rest) xs
41 findIntersecting x sets =
42 let (intersecting, disjoint) = partition (not . S.null . S.intersection x) sets
43 in if null intersecting
44 then Nothing
45 else Just (intersecting, disjoint)
46
47partition :: (a -> Bool) -> [a] -> ([a], [a])
48partition p xs = foldr select ([], []) xs
49 where
50 select x (ts, fs)
51 | p x = (x : ts, fs)
52 | otherwise = (ts, x : fs)
53
54p1 n = product $ take 3 $ map S.size $ reverse $ sortBy (comparing S.size) $ circuits n 1000
55
56p2 = circuits2
57
58main = do
59 n <- parse <$> getContents
60 print $ p1 n
61 print $ p2 n