src/task6/task6-monad.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 import Control.Monad (foldM_, ap, liftM)
   4
   5 import Data.Text (Text)
   6 import qualified Data.Text as T
   7 import qualified Data.Text.IO as TIO
   8
   9 import qualified Data.HashMap.Strict as M
  10 import Data.HashMap.Strict ((!))
  11
  12 type Name = Text
  13
  14 data Group = Group { parent :: Name
  15                    , size :: Int
  16                    } deriving (Show, Eq)
  17
  18
  19 type FriendGroups = M.HashMap Name Group
  20
  21 newtype State a = State (FriendGroups -> (FriendGroups, a))
  22
  23 instance Monad State where
  24     return x = State (\groups -> (groups, x))
  25     (State st) >>= f
  26         = State (\groups -> let
  27                             (newGroups, y) = st groups
  28                             (State trans)  = f y
  29                             in
  30                             trans newGroups
  31                 )
  32
  33 instance Applicative State where
  34     pure = return
  35     (<*>) = ap
  36
  37 instance Functor State where
  38     fmap = liftM
  39
  40
  41
  42
  43 main :: IO ()
  44 main = do
  45         friendship_text <- TIO.readFile "data/06-friendships.txt"
  46         let friendships = map enpair $ T.lines friendship_text
  47         let groups = execGroups $ mergeFriendships friendships
  48         print $ part1 groups
  49         print $ part2 groups
  50
  51
  52 part1 :: FriendGroups -> Int
  53 part1 = M.size . M.filterWithKey (\k a -> k == parent a)
  54
  55 part2 :: FriendGroups -> Int
  56 part2 = maximum . (map size) . M.elems
  57
  58
  59 enpair friendText = (this, that)
  60     where (this:that:[]) = T.words friendText
  61
  62
  63 -- run a state monad, extract the groups
  64 execGroups :: State a -> FriendGroups
  65 execGroups (State st) = fst $ st M.empty
  66
  67 -- including all the friendships is just a monadic fold
  68 mergeFriendships :: [(Name, Name)] -> State ()
  69 mergeFriendships pairs = foldM_ includeFriendshipM () pairs
  70
  71 includeFriendshipM :: () -> (Name, Name) -> State ()
  72 includeFriendshipM _ (thisPerson, thatPerson) = do
  73     include thisPerson
  74     include thatPerson
  75     thisExemplar <- exemplar thisPerson
  76     thatExemplar <- exemplar thatPerson
  77     if thisExemplar /= thatExemplar
  78     then do
  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)
  86     else
  87           return ()
  88
  89
  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
  97
  98 include :: Name -> State ()
  99 include name = State (\groups ->
 100                   if name `M.member` groups
 101                   then (groups, () )
 102                   else (M.insert name (Group {parent = name, size = 1}) groups, ())
 103                   )
 104
 105 absorb :: Group -> Group -> State ()
 106 absorb absorberG absorbedG =
 107     let absorber = parent absorberG
 108         absorbed = parent absorbedG
 109     in State (\groups ->
 110       ( M.insert absorbed ((groups!absorbed) {parent = absorber}) groups
 111       , ()
 112       )
 113     )
 114
 115 updateSize :: Group -> Int -> State ()
 116 updateSize exemplarG newSize =
 117     let name = parent exemplarG
 118     in State (\groups ->
 119       ( M.insert name ((groups!name) {size = newSize}) groups,
 120         ()
 121       )
 122     )