src/task6/task6.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 import Data.List (foldl')        -- import the strict fold
   4 -- import Data.List
   5
   6 import Data.Text (Text)
   7 import qualified Data.Text as T
   8 import qualified Data.Text.IO as TIO
   9
  10 import qualified Data.HashMap.Strict as M
  11 import Data.HashMap.Strict ((!))
  12
  13 type Name = Text
  14
  15 data Group = Group { parent :: Name
  16                    , size :: Int
  17                    } deriving (Show, Eq)
  18
  19
  20 type FriendGroups = M.HashMap Name Group
  21
  22 main :: IO ()
  23 main = do
  24         friendship_text <- TIO.readFile "data/06-friendships.txt"
  25         let friendships = map enpair $ T.lines friendship_text
  26         let groups = mergeFriendships friendships
  27         print $ part1 groups
  28         print $ part2 groups
  29
  30
  31 part1 :: FriendGroups -> Int
  32 part1 = M.size . M.filterWithKey (\k a -> k == parent a)
  33
  34 part2 :: FriendGroups -> Int
  35 part2 = maximum . (map size) . M.elems
  36
  37
  38 enpair friendText = (this, that)
  39     where (this:that:[]) = T.words friendText
  40
  41
  42 mergeFriendships :: [(Name, Name)] -> FriendGroups
  43 mergeFriendships = foldl' includeFriendship M.empty
  44
  45 exemplar :: FriendGroups -> Name -> Name
  46 exemplar groups person
  47     | person' == person = person
  48     | otherwise         = exemplar groups person'
  49     where person' = parent (groups!person)
  50
  51
  52 include :: FriendGroups -> Name -> FriendGroups
  53 include groups person =
  54     if person `M.member` groups
  55     then groups
  56     else M.insert person (Group {parent = person, size = 1}) groups
  57
  58
  59 includeFriendship :: FriendGroups -> (Name, Name) -> FriendGroups
  60 includeFriendship groups0 (thisPerson, thatPerson) =
  61     if thisExemplar == thatExemplar
  62     then groups
  63     else groups''
  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'
  75