1 {-# LANGUAGE OverloadedStrings #-}
3 import Control.Monad (foldM_, ap, liftM)
5 import Data.Text (Text)
6 import qualified Data.Text as T
7 import qualified Data.Text.IO as TIO
9 import qualified Data.HashMap.Strict as M
10 import Data.HashMap.Strict ((!))
14 data Group = Group { parent :: Name
19 type FriendGroups = M.HashMap Name Group
21 newtype State a = State (FriendGroups -> (FriendGroups, a))
23 instance Monad State where
24 return x = State (\groups -> (groups, x))
26 = State (\groups -> let
27 (newGroups, y) = st groups
33 instance Applicative State where
37 instance Functor State where
45 friendship_text <- TIO.readFile "data/06-friendships.txt"
46 let friendships = map enpair $ T.lines friendship_text
47 let groups = execGroups $ mergeFriendships friendships
52 part1 :: FriendGroups -> Int
53 part1 = M.size . M.filterWithKey (\k a -> k == parent a)
55 part2 :: FriendGroups -> Int
56 part2 = maximum . (map size) . M.elems
59 enpair friendText = (this, that)
60 where (this:that:[]) = T.words friendText
63 -- run a state monad, extract the groups
64 execGroups :: State a -> FriendGroups
65 execGroups (State st) = fst $ st M.empty
67 -- including all the friendships is just a monadic fold
68 mergeFriendships :: [(Name, Name)] -> State ()
69 mergeFriendships pairs = foldM_ includeFriendship () pairs
71 includeFriendship :: () -> (Name, Name) -> State ()
72 includeFriendship _ (thisPerson, thatPerson) = do
75 thisExemplar <- exemplar thisPerson
76 thatExemplar <- exemplar thatPerson
77 if thisExemplar /= thatExemplar
79 let thisSize = size thisExemplar
80 let thatSize = size thatExemplar
81 let (absorber, absorbed) = if thisSize > thatSize
82 then (thisExemplar, thatExemplar)
83 else (thatExemplar, thisExemplar)
84 absorb absorber absorbed
85 updateSize absorber (thisSize + thatSize)
90 exemplar :: Name -> State Group
91 exemplar name = State (exemplar' name)
92 where exemplar' person groups =
93 let person' = parent (groups!person)
94 in if person' == person
95 then (groups, groups!person)
96 else exemplar' person' groups
98 include :: Name -> State ()
99 include name = State (\groups ->
100 if name `M.member` groups
102 else (M.insert name (Group {parent = name, size = 1}) groups, ())
105 absorb :: Group -> Group -> State ()
106 absorb absorberG absorbedG =
107 let absorber = parent absorberG
108 absorbed = parent absorbedG
110 ( M.insert absorbed ((groups!absorbed) {parent = absorber}) groups
115 updateSize :: Group -> Int -> State ()
116 updateSize exemplarG newSize =
117 let name = parent exemplarG
119 ( M.insert name ((groups!name) {size = newSize}) groups,