{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -Wall #-}
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.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 {-# UNPACK #-} !Int !(Set a)
deriving Int -> RankChildren a -> ShowS
forall a. Show a => Int -> RankChildren a -> ShowS
forall a. Show a => [RankChildren a] -> ShowS
forall a. Show a => RankChildren a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RankChildren a] -> ShowS
$cshowList :: forall a. Show a => [RankChildren a] -> ShowS
show :: RankChildren a -> String
$cshow :: forall a. Show a => RankChildren a -> String
showsPrec :: Int -> RankChildren a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> RankChildren a -> ShowS
Show
data RevealDisjointSet a = RevealDisjointSet
!(Map a a)
!(Map a (RankChildren a))
deriving Int -> RevealDisjointSet a -> ShowS
forall a. Show a => Int -> RevealDisjointSet a -> ShowS
forall a. Show a => [RevealDisjointSet a] -> ShowS
forall a. Show a => RevealDisjointSet a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [RevealDisjointSet a] -> ShowS
$cshowList :: forall a. Show a => [RevealDisjointSet a] -> ShowS
show :: RevealDisjointSet a -> String
$cshow :: forall a. Show a => RevealDisjointSet a -> String
showsPrec :: Int -> RevealDisjointSet a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> RevealDisjointSet a -> ShowS
Show
showInternal :: Show a => DisjointSet a -> String
showInternal :: forall a. Show a => DisjointSet a -> String
showInternal (DisjointSet Map a a
p Map a (RankChildren a)
r) = forall a. Show a => a -> String
show (forall a. Map a a -> Map a (RankChildren a) -> RevealDisjointSet a
RevealDisjointSet Map a a
p Map a (RankChildren a)
r)
fromSets :: Ord a => [Set a] -> Maybe (DisjointSet a)
fromSets :: forall a. Ord a => [Set a] -> Maybe (DisjointSet a)
fromSets [Set a]
xs = case forall a. Ord a => [Set a] -> Maybe (Set a)
unionDistinctAll [Set a]
xs of
Maybe (Set a)
Nothing -> forall a. Maybe a
Nothing
Just Set a
_ -> forall a. a -> Maybe a
Just (forall a. Ord a => [Set a] -> DisjointSet a -> DisjointSet a
unsafeFromSets [Set a]
xs forall a. DisjointSet a
empty)
unsafeFromSets :: Ord a => [Set a] -> DisjointSet a -> DisjointSet a
unsafeFromSets :: forall a. Ord a => [Set a] -> DisjointSet a -> DisjointSet a
unsafeFromSets [Set a]
ys !ds :: DisjointSet a
ds@(DisjointSet Map a a
p Map a (RankChildren a)
r) = case [Set a]
ys of
[] -> DisjointSet a
ds
Set a
x : [Set a]
xs -> case forall a. Set a -> Maybe a
setLookupMin Set a
x of
Maybe a
Nothing -> forall a. Ord a => [Set a] -> DisjointSet a -> DisjointSet a
unsafeFromSets [Set a]
xs DisjointSet a
ds
Just a
m -> forall a. Ord a => [Set a] -> DisjointSet a -> DisjointSet a
unsafeFromSets [Set a]
xs forall a b. (a -> b) -> a -> b
$ forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet
(forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union (forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\a
_ -> a
m) Set a
x) Map a a
p)
(forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
m (forall a. Int -> Set a -> RankChildren a
RankChildren Int
0 Set a
x) Map a (RankChildren a)
r)
unionDistinctAll :: Ord a => [Set a] -> Maybe (Set a)
unionDistinctAll :: forall a. Ord a => [Set a] -> Maybe (Set a)
unionDistinctAll = forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM forall a. Ord a => Set a -> Set a -> Maybe (Set a)
unionDistinct forall a. Set a
S.empty
unionDistinct :: Ord a => Set a -> Set a -> Maybe (Set a)
unionDistinct :: forall a. Ord a => Set a -> Set a -> Maybe (Set a)
unionDistinct Set a
a Set a
b =
let s :: Set a
s = forall a. Ord a => Set a -> Set a -> Set a
S.union Set a
a Set a
b
in if forall a. Set a -> Int
S.size Set a
a forall a. Num a => a -> a -> a
+ forall a. Set a -> Int
S.size Set a
b forall a. Eq a => a -> a -> Bool
== forall a. Set a -> Int
S.size Set a
s
then forall a. a -> Maybe a
Just Set a
s
else forall a. Maybe a
Nothing
instance Ord a => Monoid (DisjointSet a) where
mappend :: DisjointSet a -> DisjointSet a -> DisjointSet a
mappend = forall a. Ord a => DisjointSet a -> DisjointSet a -> DisjointSet a
append
mempty :: DisjointSet a
mempty = forall a. DisjointSet a
empty
instance Ord a => Semigroup (DisjointSet a) where
<> :: DisjointSet a -> DisjointSet a -> DisjointSet a
(<>) = forall a. Ord a => DisjointSet a -> DisjointSet a -> DisjointSet a
append
instance Ord a => Eq (DisjointSet a) where
DisjointSet a
a == :: DisjointSet a -> DisjointSet a -> Bool
== DisjointSet a
b = forall a. Ord a => [a] -> Set a
S.fromList (forall a. DisjointSet a -> [Set a]
toSets DisjointSet a
a) forall a. Eq a => a -> a -> Bool
== forall a. Ord a => [a] -> Set a
S.fromList (forall a. DisjointSet a -> [Set a]
toSets DisjointSet a
b)
instance Ord a => Ord (DisjointSet a) where
compare :: DisjointSet a -> DisjointSet a -> Ordering
compare DisjointSet a
a DisjointSet a
b = forall a. Ord a => a -> a -> Ordering
compare (forall a. Ord a => [a] -> Set a
S.fromList (forall a. DisjointSet a -> [Set a]
toSets DisjointSet a
a)) (forall a. Ord a => [a] -> Set a
S.fromList (forall a. DisjointSet a -> [Set a]
toSets DisjointSet a
b))
instance (Show a, Ord a) => Show (DisjointSet a) where
show :: DisjointSet a -> String
show = forall a. (Show a, Ord a) => DisjointSet a -> String
showDisjointSet
showDisjointSet :: (Show a, Ord a) => DisjointSet a -> String
showDisjointSet :: forall a. (Show a, Ord a) => DisjointSet a -> String
showDisjointSet = String -> ShowS
showString String
"fromLists " forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. DisjointSet a -> [[a]]
toLists
pretty :: (Ord a, Show a) => DisjointSet a -> String
pretty :: forall a. (Ord a, Show a) => DisjointSet a -> String
pretty DisjointSet a
xs = forall a. a -> a
id
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
'{'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a -> a] -> a -> a
applyList (forall a. a -> [a] -> [a]
L.intersperse (Char -> ShowS
showChar Char
',') (forall a b. (a -> b) -> [a] -> [b]
map (\[a]
x -> Char -> ShowS
showChar Char
'{' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a -> a] -> a -> a
applyList (forall a. a -> [a] -> [a]
L.intersperse (Char -> ShowS
showChar Char
',') (forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> ShowS
shows [a]
x)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
'}') (forall a. DisjointSet a -> [[a]]
toLists DisjointSet a
xs)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
'}'
forall a b. (a -> b) -> a -> b
$ []
applyList :: [(a -> a)] -> a -> a
applyList :: forall a. [a -> a] -> a -> a
applyList [] = forall a. a -> a
id
applyList (a -> a
f : [a -> a]
fs) = a -> a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a -> a] -> a -> a
applyList [a -> a]
fs
toLists :: DisjointSet a -> [[a]]
toLists :: forall a. DisjointSet a -> [[a]]
toLists = forall a b. (a -> b) -> [a] -> [b]
map forall a. Set a -> [a]
S.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. DisjointSet a -> [Set a]
toSets
fromLists :: Ord a => [[a]] -> DisjointSet a
fromLists :: forall a. Ord a => [[a]] -> DisjointSet a
fromLists [[a]]
xs = forall a. a -> Maybe a -> a
fromMaybe forall a. DisjointSet a
empty (forall a. Ord a => [Set a] -> Maybe (DisjointSet a)
fromSets (forall a b. (a -> b) -> [a] -> [b]
map forall a. Ord a => [a] -> Set a
S.fromList [[a]]
xs))
toSets :: DisjointSet a -> [Set a]
toSets :: forall a. DisjointSet a -> [Set a]
toSets (DisjointSet Map a a
_ Map a (RankChildren a)
r) = forall a b k. (a -> b -> b) -> b -> Map k a -> b
M.foldr
(\(RankChildren Int
_ Set a
s) [Set a]
xs -> Set a
s forall a. a -> [a] -> [a]
: [Set a]
xs) [] Map a (RankChildren a)
r
union :: Ord a => a -> a -> DisjointSet a -> DisjointSet a
union :: forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
union !a
x !a
y DisjointSet a
set = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> s
execState DisjointSet a
set forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT forall a b. (a -> b) -> a -> b
$ do
a
repx <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> DisjointSet a -> (a, DisjointSet a)
lookupCompressAdd a
x
a
repy <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> DisjointSet a -> (a, DisjointSet a)
lookupCompressAdd a
y
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ a
repx forall a. Eq a => a -> a -> Bool
/= a
repy
DisjointSet Map a a
p Map a (RankChildren a)
r <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall (m :: * -> *) s. Monad m => StateT s m s
get
let RankChildren Int
rankx Set a
keysx = Map a (RankChildren a)
r forall k a. Ord k => Map k a -> k -> a
M.! a
repx
let RankChildren Int
ranky Set a
keysy = Map a (RankChildren a)
r forall k a. Ord k => Map k a -> k -> a
M.! a
repy
keys :: Set a
keys = forall a. Monoid a => a -> a -> a
mappend Set a
keysx Set a
keysy
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put forall a b. (a -> b) -> a -> b
$! case forall a. Ord a => a -> a -> Ordering
compare Int
rankx Int
ranky of
Ordering
LT -> let p' :: Map a a
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
repx a
repy Map a a
p
r' :: Map a (RankChildren a)
r' = forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
repx forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
repy (forall a. Int -> Set a -> RankChildren a
RankChildren Int
ranky Set a
keys) Map a (RankChildren a)
r
in forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p' Map a (RankChildren a)
r'
Ordering
GT -> let p' :: Map a a
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
repy a
repx Map a a
p
r' :: Map a (RankChildren a)
r' = forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
repy forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
repx (forall a. Int -> Set a -> RankChildren a
RankChildren Int
rankx Set a
keys) Map a (RankChildren a)
r
in forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p' Map a (RankChildren a)
r'
Ordering
EQ -> let p' :: Map a a
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
repx a
repy Map a a
p
r' :: Map a (RankChildren a)
r' = forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
repx forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
repy (forall a. Int -> Set a -> RankChildren a
RankChildren (Int
ranky forall a. Num a => a -> a -> a
+ Int
1) Set a
keys) Map a (RankChildren a)
r
in forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p' Map a (RankChildren a)
r'
representative :: Ord a => a -> DisjointSet a -> Maybe a
representative :: forall a. Ord a => a -> DisjointSet a -> Maybe a
representative = forall a. Ord a => a -> DisjointSet a -> Maybe a
find
equivalent :: Ord a => a -> a -> DisjointSet a -> Bool
equivalent :: forall a. Ord a => a -> a -> DisjointSet a -> Bool
equivalent a
a a
b DisjointSet a
ds = forall a. a -> Maybe a -> a
fromMaybe Bool
False forall a b. (a -> b) -> a -> b
$ do
a
x <- forall a. Ord a => a -> DisjointSet a -> Maybe a
representative a
a DisjointSet a
ds
a
y <- forall a. Ord a => a -> DisjointSet a -> Maybe a
representative a
b DisjointSet a
ds
forall a. a -> Maybe a
Just (a
x forall a. Eq a => a -> a -> Bool
== a
y)
equivalences :: Ord a => a -> DisjointSet a -> Set a
equivalences :: forall a. Ord a => a -> DisjointSet a -> Set a
equivalences a
a (DisjointSet Map a a
p Map a (RankChildren a)
r) = case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
a Map a a
p of
Maybe a
Nothing -> forall a. a -> Set a
S.singleton a
a
Just a
b -> case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall a. Ord a => a -> Map a a -> a
lookupUntilRoot a
b Map a a
p) Map a (RankChildren a)
r of
Maybe (RankChildren a)
Nothing -> forall a. HasCallStack => String -> a
error String
"Data.DisjointSet equivalences: invariant violated"
Just (RankChildren Int
_ Set a
s) -> Set a
s
lookupUntilRoot :: Ord a => a -> Map a a -> a
lookupUntilRoot :: forall a. Ord a => a -> Map a a -> a
lookupUntilRoot a
a Map a a
m = case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
a Map a a
m of
Maybe a
Nothing -> a
a
Just a
a' -> if a
a forall a. Eq a => a -> a -> Bool
== a
a'
then a
a
else forall a. Ord a => a -> Map a a -> a
lookupUntilRoot a
a' Map a a
m
sets :: DisjointSet a -> Int
sets :: forall a. DisjointSet a -> Int
sets (DisjointSet Map a a
_ Map a (RankChildren a)
r) = forall k a. Map k a -> Int
M.size Map a (RankChildren a)
r
values :: DisjointSet a -> Int
values :: forall a. DisjointSet a -> Int
values (DisjointSet Map a a
p Map a (RankChildren a)
_) = forall k a. Map k a -> Int
M.size Map a a
p
insert :: Ord a => a -> DisjointSet a -> DisjointSet a
insert :: forall a. Ord a => a -> DisjointSet a -> DisjointSet a
insert !a
x set :: DisjointSet a
set@(DisjointSet Map a a
p Map a (RankChildren a)
r) =
let (Maybe a
l, Map a a
p') = forall k a.
Ord k =>
(k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a)
M.insertLookupWithKey (\a
_ a
_ a
old -> a
old) a
x a
x Map a a
p
in case Maybe a
l of
Just a
_ -> DisjointSet a
set
Maybe a
Nothing ->
let r' :: Map a (RankChildren a)
r' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x (forall a. Int -> Set a -> RankChildren a
RankChildren Int
0 (forall a. a -> Set a
S.singleton a
x)) Map a (RankChildren a)
r
in forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p' Map a (RankChildren a)
r'
singleton :: a -> DisjointSet a
singleton :: forall a. a -> DisjointSet a
singleton !a
x =
let p :: Map a a
p = forall k a. k -> a -> Map k a
M.singleton a
x a
x
r :: Map a (RankChildren a)
r = forall k a. k -> a -> Map k a
M.singleton a
x (forall a. Int -> Set a -> RankChildren a
RankChildren Int
0 (forall a. a -> Set a
S.singleton a
x))
in forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p Map a (RankChildren a)
r
doubleton :: Ord a => a -> a -> DisjointSet a
doubleton :: forall a. Ord a => a -> a -> DisjointSet a
doubleton a
a a
b = forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
union a
a a
b forall a. DisjointSet a
empty
empty :: DisjointSet a
empty :: forall a. DisjointSet a
empty = forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet forall k a. Map k a
M.empty forall k a. Map k a
M.empty
append :: Ord a => DisjointSet a -> DisjointSet a -> DisjointSet a
append :: forall a. Ord a => DisjointSet a -> DisjointSet a -> DisjointSet a
append s1 :: DisjointSet a
s1@(DisjointSet Map a a
m1 Map a (RankChildren a)
_) s2 :: DisjointSet a
s2@(DisjointSet Map a a
m2 Map a (RankChildren a)
_) = if forall k a. Map k a -> Int
M.size Map a a
m1 forall a. Ord a => a -> a -> Bool
> forall k a. Map k a -> Int
M.size Map a a
m2
then forall a. Ord a => DisjointSet a -> Map a a -> DisjointSet a
appendParents DisjointSet a
s1 Map a a
m2
else forall a. Ord a => DisjointSet a -> Map a a -> DisjointSet a
appendParents DisjointSet a
s2 Map a a
m1
appendParents :: Ord a => DisjointSet a -> Map a a -> DisjointSet a
appendParents :: forall a. Ord a => DisjointSet a -> Map a a -> DisjointSet a
appendParents = forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' forall a b. (a -> b) -> a -> b
$ \DisjointSet a
ds a
x a
y -> if a
x forall a. Eq a => a -> a -> Bool
== a
y
then forall a. Ord a => a -> DisjointSet a -> DisjointSet a
insert a
x DisjointSet a
ds
else forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
union a
x a
y DisjointSet a
ds
singletons :: Eq a => Set a -> DisjointSet a
singletons :: forall a. Eq a => Set a -> DisjointSet a
singletons Set a
s = case forall a. Set a -> Maybe a
setLookupMin Set a
s of
Maybe a
Nothing -> forall a. DisjointSet a
empty
Just a
x ->
let p :: Map a a
p = forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\a
_ -> a
x) Set a
s
rank :: Int
rank = if forall a. Set a -> Int
S.size Set a
s forall a. Eq a => a -> a -> Bool
== Int
1 then Int
0 else Int
1
r :: Map a (RankChildren a)
r = forall k a. k -> a -> Map k a
M.singleton a
x (forall a. Int -> Set a -> RankChildren a
RankChildren Int
rank Set a
s)
in forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p Map a (RankChildren a)
r
setLookupMin :: Set a -> Maybe a
#if MIN_VERSION_containers(0,5,9)
setLookupMin :: forall a. Set a -> Maybe a
setLookupMin = forall a. Set a -> Maybe a
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' :: forall a. Ord a => a -> DisjointSet a -> (Maybe a, DisjointSet a)
representative' !a
x DisjointSet a
set =
case forall a. Ord a => a -> DisjointSet a -> Maybe a
find a
x DisjointSet a
set of
Maybe a
Nothing -> (forall a. Maybe a
Nothing, DisjointSet a
set)
Just a
rep -> let set' :: DisjointSet a
set' = forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
compress a
rep a
x DisjointSet a
set
in DisjointSet a
set' seq :: forall a b. a -> b -> b
`seq` (forall a. a -> Maybe a
Just a
rep, DisjointSet a
set')
lookupCompressAdd :: Ord a => a -> DisjointSet a -> (a, DisjointSet a)
lookupCompressAdd :: forall a. Ord a => a -> DisjointSet a -> (a, DisjointSet a)
lookupCompressAdd !a
x DisjointSet a
set =
case forall a. Ord a => a -> DisjointSet a -> Maybe a
find a
x DisjointSet a
set of
Maybe a
Nothing -> (a
x, forall a. Ord a => a -> DisjointSet a -> DisjointSet a
insert a
x DisjointSet a
set)
Just a
rep -> let set' :: DisjointSet a
set' = forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
compress a
rep a
x DisjointSet a
set
in DisjointSet a
set' seq :: forall a b. a -> b -> b
`seq` (a
rep, DisjointSet a
set')
find :: Ord a => a -> DisjointSet a -> Maybe a
find :: forall a. Ord a => a -> DisjointSet a -> Maybe a
find !a
x (DisjointSet Map a a
p Map a (RankChildren a)
_) =
do a
x' <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
x Map a a
p
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! if a
x forall a. Eq a => a -> a -> Bool
== a
x' then a
x' else a -> a
find' a
x'
where find' :: a -> a
find' a
y = let y' :: a
y' = Map a a
p forall k a. Ord k => Map k a -> k -> a
M.! a
y
in if a
y forall a. Eq a => a -> a -> Bool
== a
y' then a
y' else a -> a
find' a
y'
compress :: Ord a => a -> a -> DisjointSet a -> DisjointSet a
compress :: forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
compress !a
rep = a -> DisjointSet a -> DisjointSet a
helper
where helper :: a -> DisjointSet a -> DisjointSet a
helper !a
x set :: DisjointSet a
set@(DisjointSet Map a a
p Map a (RankChildren a)
r)
| a
x forall a. Eq a => a -> a -> Bool
== a
rep = DisjointSet a
set
| Bool
otherwise = a -> DisjointSet a -> DisjointSet a
helper a
x' DisjointSet a
set'
where x' :: a
x' = Map a a
p forall k a. Ord k => Map k a -> k -> a
M.! a
x
set' :: DisjointSet a
set' = let p' :: Map a a
p' = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x a
rep Map a a
p
in Map a a
p' seq :: forall a b. a -> b -> b
`seq` forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p' Map a (RankChildren a)
r