module Data.DisjointSet
( DisjointSet
, empty
, singleton
, singletons
, doubleton
, insert
, union
, equivalent
, sets
, values
, equivalences
, representative
, representative'
, toLists
, fromLists
, toSets
, fromSets
, pretty
, showInternal
) where
import Prelude hiding (lookup)
import Control.Monad.Trans.State.Strict
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.Class
import Control.Monad
import Data.Map (Map)
import Data.Set (Set)
import Data.Semigroup (Semigroup)
import Data.Maybe (fromMaybe)
import Data.Aeson (ToJSON(..),FromJSON(..))
import Data.Foldable (foldlM)
import qualified Data.Semigroup
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import qualified Data.List as L
data DisjointSet a = DisjointSet
!(Map a a)
!(Map a (RankChildren a))
data RankChildren a = RankChildren !Int !(Set a)
deriving Show
data RevealDisjointSet a = RevealDisjointSet
!(Map a a)
!(Map a (RankChildren a))
deriving Show
showInternal :: Show a => DisjointSet a -> String
showInternal (DisjointSet p r) = show (RevealDisjointSet p r)
instance ToJSON a => ToJSON (DisjointSet a) where
toJSON = toJSON . toSets
instance (Ord a, FromJSON a) => FromJSON (DisjointSet a) where
parseJSON x = do
theSets <- parseJSON x
case fromSets theSets of
Nothing -> fail "the sets comprising the DisjointSet were not distinct"
Just s -> return s
fromSets :: Ord a => [Set a] -> Maybe (DisjointSet a)
fromSets xs = case unionDistinctAll xs of
Nothing -> Nothing
Just _ -> Just (unsafeFromSets xs empty)
unsafeFromSets :: Ord a => [Set a] -> DisjointSet a -> DisjointSet a
unsafeFromSets ys !ds@(DisjointSet p r) = case ys of
[] -> ds
x : xs -> case setLookupMin x of
Nothing -> unsafeFromSets xs ds
Just m -> unsafeFromSets xs $ DisjointSet
(M.union (M.fromSet (\_ -> m) x) p)
(M.insert m (RankChildren 0 x) r)
unionDistinctAll :: Ord a => [Set a] -> Maybe (Set a)
unionDistinctAll = foldlM unionDistinct S.empty
unionDistinct :: Ord a => Set a -> Set a -> Maybe (Set a)
unionDistinct a b =
let s = S.union a b
in if S.size a + S.size b == S.size s
then Just s
else Nothing
instance Ord a => Monoid (DisjointSet a) where
mappend = append
mempty = empty
instance Ord a => Semigroup (DisjointSet a) where
(<>) = append
instance Ord a => Eq (DisjointSet a) where
a == b = S.fromList (toSets a) == S.fromList (toSets b)
instance Ord a => Ord (DisjointSet a) where
compare a b = compare (S.fromList (toSets a)) (S.fromList (toSets b))
instance (Show a, Ord a) => Show (DisjointSet a) where
show = showDisjointSet
showDisjointSet :: (Show a, Ord a) => DisjointSet a -> String
showDisjointSet = showString "fromLists " . show . toLists
pretty :: (Ord a, Show a) => DisjointSet a -> String
pretty xs = id
. showChar '{'
. applyList (L.intersperse (showChar ',') (map (\x -> showChar '{' . applyList (L.intersperse (showChar ',') (map shows x)) . showChar '}') (toLists xs)))
. showChar '}'
$ []
applyList :: [(a -> a)] -> a -> a
applyList [] = id
applyList (f : fs) = f . applyList fs
toLists :: DisjointSet a -> [[a]]
toLists = map S.toList . toSets
fromLists :: Ord a => [[a]] -> DisjointSet a
fromLists xs = fromMaybe empty (fromSets (map S.fromList xs))
toSets :: DisjointSet a -> [Set a]
toSets (DisjointSet _ r) = M.foldr
(\(RankChildren _ s) xs -> s : xs) [] r
union :: Ord a => a -> a -> DisjointSet a -> DisjointSet a
union !x !y set = flip execState set $ runMaybeT $ do
repx <- lift $ state $ lookupCompressAdd x
repy <- lift $ state $ lookupCompressAdd y
guard $ repx /= repy
DisjointSet p r <- lift get
let RankChildren rankx keysx = r M.! repx
let RankChildren ranky keysy = r M.! repy
keys = mappend keysx keysy
lift $ put $! case compare rankx ranky of
LT -> let p' = M.insert repx repy p
r' = M.delete repx $! M.insert repy (RankChildren ranky keys) r
in DisjointSet p' r'
GT -> let p' = M.insert repy repx p
r' = M.delete repy $! M.insert repx (RankChildren rankx keys) r
in DisjointSet p' r'
EQ -> let p' = M.insert repx repy p
r' = M.delete repx $! M.insert repy (RankChildren (ranky + 1) keys) r
in DisjointSet p' r'
representative :: Ord a => a -> DisjointSet a -> Maybe a
representative = find
equivalent :: Ord a => a -> a -> DisjointSet a -> Bool
equivalent a b ds = fromMaybe False $ do
x <- representative a ds
y <- representative b ds
Just (x == y)
equivalences :: Ord a => a -> DisjointSet a -> Set a
equivalences a (DisjointSet p r) = case M.lookup a p of
Nothing -> S.singleton a
Just b -> case M.lookup (lookupUntilRoot b p) r of
Nothing -> error "Data.DisjointSet equivalences: invariant violated"
Just (RankChildren _ s) -> s
lookupUntilRoot :: Ord a => a -> Map a a -> a
lookupUntilRoot a m = case M.lookup a m of
Nothing -> a
Just a' -> if a == a'
then a
else lookupUntilRoot a' m
sets :: DisjointSet a -> Int
sets (DisjointSet _ r) = M.size r
values :: DisjointSet a -> Int
values (DisjointSet p _) = M.size p
insert :: Ord a => a -> DisjointSet a -> DisjointSet a
insert !x set@(DisjointSet p r) =
let (l, p') = M.insertLookupWithKey (\_ _ old -> old) x x p
in case l of
Just _ -> set
Nothing ->
let r' = M.insert x (RankChildren 0 (S.singleton x)) r
in DisjointSet p' r'
singleton :: a -> DisjointSet a
singleton !x =
let p = M.singleton x x
r = M.singleton x (RankChildren 0 (S.singleton x))
in DisjointSet p r
doubleton :: Ord a => a -> a -> DisjointSet a
doubleton a b = union a b empty
empty :: DisjointSet a
empty = DisjointSet M.empty M.empty
append :: Ord a => DisjointSet a -> DisjointSet a -> DisjointSet a
append s1@(DisjointSet m1 _) s2@(DisjointSet m2 _) = if M.size m1 > M.size m2
then appendParents s1 m2
else appendParents s2 m1
appendParents :: Ord a => DisjointSet a -> Map a a -> DisjointSet a
appendParents = M.foldlWithKey' $ \ds x y -> if x == y
then insert x ds
else union x y ds
singletons :: Eq a => Set a -> DisjointSet a
singletons s = case setLookupMin s of
Nothing -> empty
Just x ->
let p = M.fromSet (\_ -> x) s
rank = if S.size s == 1 then 0 else 1
r = M.singleton x (RankChildren rank s)
in DisjointSet p r
setLookupMin :: Set a -> Maybe a
#if MIN_VERSION_containers(0,5,9)
setLookupMin = S.lookupMin
#else
setLookupMin s = if S.size s > 0 then Just (S.findMin s) else Nothing
#endif
representative' :: Ord a => a -> DisjointSet a -> (Maybe a, DisjointSet a)
representative' !x set =
case find x set of
Nothing -> (Nothing, set)
Just rep -> let set' = compress rep x set
in set' `seq` (Just rep, set')
lookupCompressAdd :: Ord a => a -> DisjointSet a -> (a, DisjointSet a)
lookupCompressAdd !x set =
case find x set of
Nothing -> (x, insert x set)
Just rep -> let set' = compress rep x set
in set' `seq` (rep, set')
find :: Ord a => a -> DisjointSet a -> Maybe a
find !x (DisjointSet p _) =
do x' <- M.lookup x p
return $! if x == x' then x' else find' x'
where find' y = let y' = p M.! y
in if y == y' then y' else find' y'
compress :: Ord a => a -> a -> DisjointSet a -> DisjointSet a
compress !rep = helper
where helper !x set@(DisjointSet p r)
| x == rep = set
| otherwise = helper x' set'
where x' = p M.! x
set' = let p' = M.insert x rep p
in p' `seq` DisjointSet p' r