1 {-# LANGUAGE OverloadedStrings #-}
3 import Data.List (foldl') -- import the strict fold
6 import Data.Text (Text)
7 import qualified Data.Text as T
8 import qualified Data.Text.IO as TIO
10 import qualified Data.HashMap.Strict as M
11 import Data.HashMap.Strict ((!))
15 data Group = Group { parent :: Name
20 type FriendGroups = M.HashMap Name Group
24 friendship_text <- TIO.readFile "data/06-friendships.txt"
25 let friendships = map enpair $ T.lines friendship_text
26 let groups = mergeFriendships friendships
31 part1 :: FriendGroups -> Int
32 part1 = M.size . M.filterWithKey (\k a -> k == parent a)
34 part2 :: FriendGroups -> Int
35 part2 = maximum . (map size) . M.elems
38 enpair friendText = (this, that)
39 where (this:that:[]) = T.words friendText
42 mergeFriendships :: [(Name, Name)] -> FriendGroups
43 mergeFriendships = foldl' includeFriendship M.empty
45 exemplar :: FriendGroups -> Name -> Name
46 exemplar groups person
47 | person' == person = person
48 | otherwise = exemplar groups person'
49 where person' = parent (groups!person)
52 include :: FriendGroups -> Name -> FriendGroups
53 include groups person =
54 if person `M.member` groups
56 else M.insert person (Group {parent = person, size = 1}) groups
59 includeFriendship :: FriendGroups -> (Name, Name) -> FriendGroups
60 includeFriendship groups0 (thisPerson, thatPerson) =
61 if thisExemplar == thatExemplar
64 where groups1 = include groups0 thisPerson
65 groups = include groups1 thatPerson
66 thisExemplar = exemplar groups thisPerson
67 thatExemplar = exemplar groups thatPerson
68 thisSize = size $ groups!thisExemplar
69 thatSize = size $ groups!thatExemplar
70 (absorber, absorbed) = if thisSize > thatSize
71 then (thisExemplar, thatExemplar)
72 else (thatExemplar, thisExemplar)
73 groups' = M.insert absorbed ((groups!absorbed) {parent = absorber}) groups
74 groups'' = M.insert absorber ((groups!absorber) {size = thisSize + thatSize}) groups'