25 lines
816 B
Haskell
25 lines
816 B
Haskell
{-# LANGUAGE TupleSections #-}
|
|
|
|
module Rextra.Util
|
|
( connectedElements
|
|
, groupByFirst
|
|
) where
|
|
|
|
import Control.Monad
|
|
import Control.Monad.Trans.State
|
|
import qualified Data.Map as Map
|
|
import qualified Data.Set as Set
|
|
|
|
explore :: (Ord n) => (n -> Set.Set n) -> n -> State (Set.Set n) ()
|
|
explore trans node = do
|
|
visited <- get
|
|
unless (node `Set.member` visited) $ do
|
|
modify (Set.insert node)
|
|
mapM_ (explore trans) . Set.toList $ trans node
|
|
|
|
connectedElements :: (Ord n) => (n -> Set.Set n) -> Set.Set n -> Set.Set n
|
|
connectedElements trans startingNodes =
|
|
flip execState Set.empty . mapM (explore trans) $ Set.toList startingNodes
|
|
|
|
groupByFirst :: (Ord a, Ord b) => [(a, b)] -> Map.Map a (Set.Set b)
|
|
groupByFirst = Map.fromListWith mappend . map (\(a, b) -> (a, Set.singleton b))
|