{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -Wall #-}

{- |
Persistent disjoint-sets. Disjoint-sets are a set of elements
with equivalence relations defined between elements, i.e.
two elements may be members of the same equivalence set.
The type in this module can be roughly understood as:

> DisjointSet a ≈ Set (Set a)

This library provides the fundamental operations classically
known as @union@, @find@, and @makeSet@. It also offers
novelties like a 'Monoid' instance for disjoint sets
and conversion functions for interoperating with lists.
See the tutorial at the bottom of this page for an example
of how to use this library.
-}
module Data.DisjointSet
  ( DisjointSet

    -- * Construction
  , empty
  , singleton
  , singletons
  , doubleton
  , insert
  , union

    -- * Query
  , equivalent
  , sets
  , values
  , equivalences
  , representative
  , representative'

    -- * Conversion
  , toLists
  , fromLists
  , toSets
  , fromSets
  , pretty
  , showInternal

    -- * Tutorial
    -- $tutorial
  ) where

import Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.State.Strict
import Prelude hiding (lookup)

import Data.Foldable (foldlM)
import qualified Data.List as L
import Data.Map (Map)
import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe)
import Data.Set (Set)
import qualified Data.Set as S

data DisjointSet a
  = DisjointSet
      !(Map a a) -- parents
      !(Map a (RankChildren a)) -- ranks

data RankChildren a = RankChildren {-# UNPACK #-} !Int !(Set a)
  deriving (Int -> RankChildren a -> ShowS
[RankChildren a] -> ShowS
RankChildren a -> String
(Int -> RankChildren a -> ShowS)
-> (RankChildren a -> String)
-> ([RankChildren a] -> ShowS)
-> Show (RankChildren a)
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
$cshowsPrec :: forall a. Show a => Int -> RankChildren a -> ShowS
showsPrec :: Int -> RankChildren a -> ShowS
$cshow :: forall a. Show a => RankChildren a -> String
show :: RankChildren a -> String
$cshowList :: forall a. Show a => [RankChildren a] -> ShowS
showList :: [RankChildren a] -> ShowS
Show)

data RevealDisjointSet a
  = RevealDisjointSet
      !(Map a a)
      !(Map a (RankChildren a))
  deriving (Int -> RevealDisjointSet a -> ShowS
[RevealDisjointSet a] -> ShowS
RevealDisjointSet a -> String
(Int -> RevealDisjointSet a -> ShowS)
-> (RevealDisjointSet a -> String)
-> ([RevealDisjointSet a] -> ShowS)
-> Show (RevealDisjointSet a)
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
$cshowsPrec :: forall a. Show a => Int -> RevealDisjointSet a -> ShowS
showsPrec :: Int -> RevealDisjointSet a -> ShowS
$cshow :: forall a. Show a => RevealDisjointSet a -> String
show :: RevealDisjointSet a -> String
$cshowList :: forall a. Show a => [RevealDisjointSet a] -> ShowS
showList :: [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) = RevealDisjointSet a -> String
forall a. Show a => a -> String
show (Map a a -> Map a (RankChildren a) -> RevealDisjointSet a
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 [Set a] -> Maybe (Set a)
forall a. Ord a => [Set a] -> Maybe (Set a)
unionDistinctAll [Set a]
xs of
  Maybe (Set a)
Nothing -> Maybe (DisjointSet a)
forall a. Maybe a
Nothing
  Just Set a
_ -> DisjointSet a -> Maybe (DisjointSet a)
forall a. a -> Maybe a
Just ([Set a] -> DisjointSet a -> DisjointSet a
forall a. Ord a => [Set a] -> DisjointSet a -> DisjointSet a
unsafeFromSets [Set a]
xs DisjointSet a
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 Set a -> Maybe a
forall a. Set a -> Maybe a
setLookupMin Set a
x of
    Maybe a
Nothing -> [Set a] -> DisjointSet a -> DisjointSet a
forall a. Ord a => [Set a] -> DisjointSet a -> DisjointSet a
unsafeFromSets [Set a]
xs DisjointSet a
ds
    Just a
m ->
      [Set a] -> DisjointSet a -> DisjointSet a
forall a. Ord a => [Set a] -> DisjointSet a -> DisjointSet a
unsafeFromSets [Set a]
xs (DisjointSet a -> DisjointSet a) -> DisjointSet a -> DisjointSet a
forall a b. (a -> b) -> a -> b
$
        Map a a -> Map a (RankChildren a) -> DisjointSet a
forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet
          (Map a a -> Map a a -> Map a a
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union ((a -> a) -> Set a -> Map a a
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\a
_ -> a
m) Set a
x) Map a a
p)
          (a
-> RankChildren a
-> Map a (RankChildren a)
-> Map a (RankChildren a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
m (Int -> Set a -> RankChildren a
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 = (Set a -> Set a -> Maybe (Set a))
-> Set a -> [Set a] -> Maybe (Set a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM Set a -> Set a -> Maybe (Set a)
forall a. Ord a => Set a -> Set a -> Maybe (Set a)
unionDistinct Set a
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 = Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.union Set a
a Set a
b
   in if Set a -> Int
forall a. Set a -> Int
S.size Set a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Set a -> Int
forall a. Set a -> Int
S.size Set a
b Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Set a -> Int
forall a. Set a -> Int
S.size Set a
s
        then Set a -> Maybe (Set a)
forall a. a -> Maybe a
Just Set a
s
        else Maybe (Set a)
forall a. Maybe a
Nothing

instance (Ord a) => Monoid (DisjointSet a) where
  mappend :: DisjointSet a -> DisjointSet a -> DisjointSet a
mappend = DisjointSet a -> DisjointSet a -> DisjointSet a
forall a. Semigroup a => a -> a -> a
(<>)
  mempty :: DisjointSet a
mempty = DisjointSet a
forall a. DisjointSet a
empty

instance (Ord a) => Semigroup (DisjointSet a) where
  <> :: DisjointSet a -> DisjointSet a -> DisjointSet a
(<>) = 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 = [Set a] -> Set (Set a)
forall a. Ord a => [a] -> Set a
S.fromList (DisjointSet a -> [Set a]
forall a. DisjointSet a -> [Set a]
toSets DisjointSet a
a) Set (Set a) -> Set (Set a) -> Bool
forall a. Eq a => a -> a -> Bool
== [Set a] -> Set (Set a)
forall a. Ord a => [a] -> Set a
S.fromList (DisjointSet a -> [Set a]
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 = Set (Set a) -> Set (Set a) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ([Set a] -> Set (Set a)
forall a. Ord a => [a] -> Set a
S.fromList (DisjointSet a -> [Set a]
forall a. DisjointSet a -> [Set a]
toSets DisjointSet a
a)) ([Set a] -> Set (Set a)
forall a. Ord a => [a] -> Set a
S.fromList (DisjointSet a -> [Set a]
forall a. DisjointSet a -> [Set a]
toSets DisjointSet a
b))

instance (Show a, Ord a) => Show (DisjointSet a) where
  show :: DisjointSet a -> String
show = DisjointSet a -> String
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 " ShowS -> (DisjointSet a -> String) -> DisjointSet a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> String
forall a. Show a => a -> String
show ([[a]] -> String)
-> (DisjointSet a -> [[a]]) -> DisjointSet a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DisjointSet a -> [[a]]
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 =
  Char -> ShowS
showChar Char
'{'
    ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ShowS] -> ShowS
forall a. [a -> a] -> a -> a
applyList (ShowS -> [ShowS] -> [ShowS]
forall a. a -> [a] -> [a]
L.intersperse (Char -> ShowS
showChar Char
',') (([a] -> ShowS) -> [[a]] -> [ShowS]
forall a b. (a -> b) -> [a] -> [b]
map (\[a]
x -> Char -> ShowS
showChar Char
'{' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ShowS] -> ShowS
forall a. [a -> a] -> a -> a
applyList (ShowS -> [ShowS] -> [ShowS]
forall a. a -> [a] -> [a]
L.intersperse (Char -> ShowS
showChar Char
',') ((a -> ShowS) -> [a] -> [ShowS]
forall a b. (a -> b) -> [a] -> [b]
map a -> ShowS
forall a. Show a => a -> ShowS
shows [a]
x)) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
'}') (DisjointSet a -> [[a]]
forall a. DisjointSet a -> [[a]]
toLists DisjointSet a
xs)))
    ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
'}'
    ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ []

applyList :: [a -> a] -> a -> a
applyList :: forall a. [a -> a] -> a -> a
applyList [] = a -> a
forall a. a -> a
id
applyList (a -> a
f : [a -> a]
fs) = a -> a
f (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a -> a] -> a -> a
forall a. [a -> a] -> a -> a
applyList [a -> a]
fs

toLists :: DisjointSet a -> [[a]]
toLists :: forall a. DisjointSet a -> [[a]]
toLists = (Set a -> [a]) -> [Set a] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map Set a -> [a]
forall a. Set a -> [a]
S.toList ([Set a] -> [[a]])
-> (DisjointSet a -> [Set a]) -> DisjointSet a -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DisjointSet a -> [Set a]
forall a. DisjointSet a -> [Set a]
toSets

-- this definition is pretty awful. Come up with something that
-- behaves a little more reasonably in the presence of failure.
fromLists :: (Ord a) => [[a]] -> DisjointSet a
fromLists :: forall a. Ord a => [[a]] -> DisjointSet a
fromLists [[a]]
xs = DisjointSet a -> Maybe (DisjointSet a) -> DisjointSet a
forall a. a -> Maybe a -> a
fromMaybe DisjointSet a
forall a. DisjointSet a
empty ([Set a] -> Maybe (DisjointSet a)
forall a. Ord a => [Set a] -> Maybe (DisjointSet a)
fromSets (([a] -> Set a) -> [[a]] -> [Set a]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> Set a
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) =
  (RankChildren a -> [Set a] -> [Set a])
-> [Set a] -> Map a (RankChildren a) -> [Set a]
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 Set a -> [Set a] -> [Set a]
forall a. a -> [a] -> [a]
: [Set a]
xs)
    []
    Map a (RankChildren a)
r

{- |
Create an equivalence relation between x and y. If either x or y
are not already is the disjoint set, they are first created
as singletons.
-}
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 = (State (DisjointSet a) (Maybe ())
 -> DisjointSet a -> DisjointSet a)
-> DisjointSet a
-> State (DisjointSet a) (Maybe ())
-> DisjointSet a
forall a b c. (a -> b -> c) -> b -> a -> c
flip State (DisjointSet a) (Maybe ()) -> DisjointSet a -> DisjointSet a
forall s a. State s a -> s -> s
execState DisjointSet a
set (State (DisjointSet a) (Maybe ()) -> DisjointSet a)
-> State (DisjointSet a) (Maybe ()) -> DisjointSet a
forall a b. (a -> b) -> a -> b
$ MaybeT (StateT (DisjointSet a) Identity) ()
-> State (DisjointSet a) (Maybe ())
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT (StateT (DisjointSet a) Identity) ()
 -> State (DisjointSet a) (Maybe ()))
-> MaybeT (StateT (DisjointSet a) Identity) ()
-> State (DisjointSet a) (Maybe ())
forall a b. (a -> b) -> a -> b
$ do
  a
repx <- StateT (DisjointSet a) Identity a
-> MaybeT (StateT (DisjointSet a) Identity) a
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointSet a) Identity a
 -> MaybeT (StateT (DisjointSet a) Identity) a)
-> StateT (DisjointSet a) Identity a
-> MaybeT (StateT (DisjointSet a) Identity) a
forall a b. (a -> b) -> a -> b
$ (DisjointSet a -> (a, DisjointSet a))
-> StateT (DisjointSet a) Identity a
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((DisjointSet a -> (a, DisjointSet a))
 -> StateT (DisjointSet a) Identity a)
-> (DisjointSet a -> (a, DisjointSet a))
-> StateT (DisjointSet a) Identity a
forall a b. (a -> b) -> a -> b
$ a -> DisjointSet a -> (a, DisjointSet a)
forall a. Ord a => a -> DisjointSet a -> (a, DisjointSet a)
lookupCompressAdd a
x
  a
repy <- StateT (DisjointSet a) Identity a
-> MaybeT (StateT (DisjointSet a) Identity) a
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointSet a) Identity a
 -> MaybeT (StateT (DisjointSet a) Identity) a)
-> StateT (DisjointSet a) Identity a
-> MaybeT (StateT (DisjointSet a) Identity) a
forall a b. (a -> b) -> a -> b
$ (DisjointSet a -> (a, DisjointSet a))
-> StateT (DisjointSet a) Identity a
forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
state ((DisjointSet a -> (a, DisjointSet a))
 -> StateT (DisjointSet a) Identity a)
-> (DisjointSet a -> (a, DisjointSet a))
-> StateT (DisjointSet a) Identity a
forall a b. (a -> b) -> a -> b
$ a -> DisjointSet a -> (a, DisjointSet a)
forall a. Ord a => a -> DisjointSet a -> (a, DisjointSet a)
lookupCompressAdd a
y
  Bool -> MaybeT (StateT (DisjointSet a) Identity) ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> MaybeT (StateT (DisjointSet a) Identity) ())
-> Bool -> MaybeT (StateT (DisjointSet a) Identity) ()
forall a b. (a -> b) -> a -> b
$ a
repx a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
repy
  DisjointSet Map a a
p Map a (RankChildren a)
r <- StateT (DisjointSet a) Identity (DisjointSet a)
-> MaybeT (StateT (DisjointSet a) Identity) (DisjointSet a)
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift StateT (DisjointSet a) Identity (DisjointSet a)
forall (m :: * -> *) s. Monad m => StateT s m s
get
  let RankChildren Int
rankx Set a
keysx = Map a (RankChildren a)
r Map a (RankChildren a) -> a -> RankChildren a
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 Map a (RankChildren a) -> a -> RankChildren a
forall k a. Ord k => Map k a -> k -> a
M.! a
repy
      keys :: Set a
keys = Set a -> Set a -> Set a
forall a. Monoid a => a -> a -> a
mappend Set a
keysx Set a
keysy
  StateT (DisjointSet a) Identity ()
-> MaybeT (StateT (DisjointSet a) Identity) ()
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT (DisjointSet a) Identity ()
 -> MaybeT (StateT (DisjointSet a) Identity) ())
-> StateT (DisjointSet a) Identity ()
-> MaybeT (StateT (DisjointSet a) Identity) ()
forall a b. (a -> b) -> a -> b
$ DisjointSet a -> StateT (DisjointSet a) Identity ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put (DisjointSet a -> StateT (DisjointSet a) Identity ())
-> DisjointSet a -> StateT (DisjointSet a) Identity ()
forall a b. (a -> b) -> a -> b
$! case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
rankx Int
ranky of
    Ordering
LT ->
      let p' :: Map a a
p' = a -> a -> Map a a -> Map a a
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' = a -> Map a (RankChildren a) -> Map a (RankChildren a)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
repx (Map a (RankChildren a) -> Map a (RankChildren a))
-> Map a (RankChildren a) -> Map a (RankChildren a)
forall a b. (a -> b) -> a -> b
$! a
-> RankChildren a
-> Map a (RankChildren a)
-> Map a (RankChildren a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
repy (Int -> Set a -> RankChildren a
forall a. Int -> Set a -> RankChildren a
RankChildren Int
ranky Set a
keys) Map a (RankChildren a)
r
       in Map a a -> Map a (RankChildren a) -> DisjointSet a
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' = a -> a -> Map a a -> Map a a
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' = a -> Map a (RankChildren a) -> Map a (RankChildren a)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
repy (Map a (RankChildren a) -> Map a (RankChildren a))
-> Map a (RankChildren a) -> Map a (RankChildren a)
forall a b. (a -> b) -> a -> b
$! a
-> RankChildren a
-> Map a (RankChildren a)
-> Map a (RankChildren a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
repx (Int -> Set a -> RankChildren a
forall a. Int -> Set a -> RankChildren a
RankChildren Int
rankx Set a
keys) Map a (RankChildren a)
r
       in Map a a -> Map a (RankChildren a) -> DisjointSet a
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' = a -> a -> Map a a -> Map a a
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' = a -> Map a (RankChildren a) -> Map a (RankChildren a)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
repx (Map a (RankChildren a) -> Map a (RankChildren a))
-> Map a (RankChildren a) -> Map a (RankChildren a)
forall a b. (a -> b) -> a -> b
$! a
-> RankChildren a
-> Map a (RankChildren a)
-> Map a (RankChildren a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
repy (Int -> Set a -> RankChildren a
forall a. Int -> Set a -> RankChildren a
RankChildren (Int
ranky Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Set a
keys) Map a (RankChildren a)
r
       in Map a a -> Map a (RankChildren a) -> DisjointSet a
forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p' Map a (RankChildren a)
r'

{- |
Find the set representative for this input.
-}
representative :: (Ord a) => a -> DisjointSet a -> Maybe a
representative :: forall a. Ord a => a -> DisjointSet a -> Maybe a
representative = a -> DisjointSet a -> Maybe a
forall a. Ord a => a -> DisjointSet a -> Maybe a
find

-- | Decides whether the two values belong to the same set
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 = Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
False (Maybe Bool -> Bool) -> Maybe Bool -> Bool
forall a b. (a -> b) -> a -> b
$ do
  a
x <- a -> DisjointSet a -> Maybe a
forall a. Ord a => a -> DisjointSet a -> Maybe a
representative a
a DisjointSet a
ds
  a
y <- a -> DisjointSet a -> Maybe a
forall a. Ord a => a -> DisjointSet a -> Maybe a
representative a
b DisjointSet a
ds
  Bool -> Maybe Bool
forall a. a -> Maybe a
Just (a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y)

{- | All elements the are considered equal to the value. In the event
    that the element does not exist, a singleton set will be returned.
-}
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 a -> Map a a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
a Map a a
p of
  Maybe a
Nothing -> a -> Set a
forall a. a -> Set a
S.singleton a
a
  Just a
b -> case a -> Map a (RankChildren a) -> Maybe (RankChildren a)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (a -> Map a a -> a
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 -> String -> Set a
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 a -> Map a a -> Maybe a
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 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
a'
      then a
a
      else a -> Map a a -> a
forall a. Ord a => a -> Map a a -> a
lookupUntilRoot a
a' Map a a
m

-- | Count the number of disjoint sets
sets :: DisjointSet a -> Int
sets :: forall a. DisjointSet a -> Int
sets (DisjointSet Map a a
_ Map a (RankChildren a)
r) = Map a (RankChildren a) -> Int
forall k a. Map k a -> Int
M.size Map a (RankChildren a)
r

-- | Count the total number of values contained by the disjoint sets
values :: DisjointSet a -> Int
values :: forall a. DisjointSet a -> Int
values (DisjointSet Map a a
p Map a (RankChildren a)
_) = Map a a -> Int
forall k a. Map k a -> Int
M.size Map a a
p

{- | Insert x into the disjoint set.  If it is already a member,
    then do nothing, otherwise x has no equivalence relations.
    O(logn).
-}
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') = (a -> a -> a -> a) -> a -> a -> Map a a -> (Maybe a, Map a a)
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' = a
-> RankChildren a
-> Map a (RankChildren a)
-> Map a (RankChildren a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x (Int -> Set a -> RankChildren a
forall a. Int -> Set a -> RankChildren a
RankChildren Int
0 (a -> Set a
forall a. a -> Set a
S.singleton a
x)) Map a (RankChildren a)
r
           in Map a a -> Map a (RankChildren a) -> DisjointSet a
forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p' Map a (RankChildren a)
r'

-- | Create a disjoint set with one member. O(1).
singleton :: a -> DisjointSet a
singleton :: forall a. a -> DisjointSet a
singleton !a
x =
  let p :: Map a a
p = a -> a -> Map a a
forall k a. k -> a -> Map k a
M.singleton a
x a
x
      r :: Map a (RankChildren a)
r = a -> RankChildren a -> Map a (RankChildren a)
forall k a. k -> a -> Map k a
M.singleton a
x (Int -> Set a -> RankChildren a
forall a. Int -> Set a -> RankChildren a
RankChildren Int
0 (a -> Set a
forall a. a -> Set a
S.singleton a
x))
   in Map a a -> Map a (RankChildren a) -> DisjointSet a
forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p Map a (RankChildren a)
r

-- | Create a disjoint set with a single set containing two members
doubleton :: (Ord a) => a -> a -> DisjointSet a
doubleton :: forall a. Ord a => a -> a -> DisjointSet a
doubleton a
a a
b = a -> a -> DisjointSet a -> DisjointSet a
forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
union a
a a
b DisjointSet a
forall a. DisjointSet a
empty

-- doubleton could be more efficient

-- | The empty set of disjoint sets.
empty :: DisjointSet a
empty :: forall a. DisjointSet a
empty = Map a a -> Map a (RankChildren a) -> DisjointSet a
forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
forall k a. Map k a
M.empty Map a (RankChildren a)
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 Map a a -> Int
forall k a. Map k a -> Int
M.size Map a a
m1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Map a a -> Int
forall k a. Map k a -> Int
M.size Map a a
m2
    then DisjointSet a -> Map a a -> DisjointSet a
forall a. Ord a => DisjointSet a -> Map a a -> DisjointSet a
appendParents DisjointSet a
s1 Map a a
m2
    else DisjointSet a -> Map a a -> DisjointSet a
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 = (DisjointSet a -> a -> a -> DisjointSet a)
-> DisjointSet a -> Map a a -> DisjointSet a
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' ((DisjointSet a -> a -> a -> DisjointSet a)
 -> DisjointSet a -> Map a a -> DisjointSet a)
-> (DisjointSet a -> a -> a -> DisjointSet a)
-> DisjointSet a
-> Map a a
-> DisjointSet a
forall a b. (a -> b) -> a -> b
$ \DisjointSet a
ds a
x a
y ->
  if a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y
    then a -> DisjointSet a -> DisjointSet a
forall a. Ord a => a -> DisjointSet a -> DisjointSet a
insert a
x DisjointSet a
ds
    else a -> a -> DisjointSet a -> DisjointSet a
forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
union a
x a
y DisjointSet a
ds

-- | Create a disjoint set where all members are equal.
singletons :: (Eq a) => Set a -> DisjointSet a
singletons :: forall a. Eq a => Set a -> DisjointSet a
singletons Set a
s = case Set a -> Maybe a
forall a. Set a -> Maybe a
setLookupMin Set a
s of
  Maybe a
Nothing -> DisjointSet a
forall a. DisjointSet a
empty
  Just a
x ->
    let p :: Map a a
p = (a -> a) -> Set a -> Map a a
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\a
_ -> a
x) Set a
s
        rank :: Int
rank = if Set a -> Int
forall a. Set a -> Int
S.size Set a
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 then Int
0 else Int
1
        r :: Map a (RankChildren a)
r = a -> RankChildren a -> Map a (RankChildren a)
forall k a. k -> a -> Map k a
M.singleton a
x (Int -> Set a -> RankChildren a
forall a. Int -> Set a -> RankChildren a
RankChildren Int
rank Set a
s)
     in Map a a -> Map a (RankChildren a) -> DisjointSet a
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 = Set a -> Maybe a
forall a. Set a -> Maybe a
S.lookupMin
#else
setLookupMin s = if S.size s > 0 then Just (S.findMin s) else Nothing
#endif

{- |
Find the set representative for this input. Returns a new disjoint
set in which the path has been compressed.
-}
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 a -> DisjointSet a -> Maybe a
forall a. Ord a => a -> DisjointSet a -> Maybe a
find a
x DisjointSet a
set of
    Maybe a
Nothing -> (Maybe a
forall a. Maybe a
Nothing, DisjointSet a
set)
    Just a
rep ->
      let set' :: DisjointSet a
set' = a -> a -> DisjointSet a -> DisjointSet a
forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
compress a
rep a
x DisjointSet a
set
       in DisjointSet a
set' DisjointSet a
-> (Maybe a, DisjointSet a) -> (Maybe a, DisjointSet a)
forall a b. a -> b -> b
`seq` (a -> Maybe a
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 a -> DisjointSet a -> Maybe a
forall a. Ord a => a -> DisjointSet a -> Maybe a
find a
x DisjointSet a
set of
    Maybe a
Nothing -> (a
x, a -> DisjointSet a -> DisjointSet a
forall a. Ord a => a -> DisjointSet a -> DisjointSet a
insert a
x DisjointSet a
set)
    Just a
rep ->
      let set' :: DisjointSet a
set' = a -> a -> DisjointSet a -> DisjointSet a
forall a. Ord a => a -> a -> DisjointSet a -> DisjointSet a
compress a
rep a
x DisjointSet a
set
       in DisjointSet a
set' DisjointSet a -> (a, DisjointSet a) -> (a, DisjointSet a)
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' <- a -> Map a a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
x Map a a
p
    a -> Maybe a
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! if a
x a -> a -> Bool
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 Map a a -> a -> a
forall k a. Ord k => Map k a -> k -> a
M.! a
y
     in if a
y a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y' then a
y' else a -> a
find' a
y'

-- TODO: make this smarter about recreating the parents Map.
-- Currently, it will recreate it more often than needed.
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 a -> a -> Bool
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 Map a a -> a -> a
forall k a. Ord k => Map k a -> k -> a
M.! a
x
    set' :: DisjointSet a
set' =
      let p' :: Map a a
p' = a -> a -> Map a a -> Map a a
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' Map a a -> DisjointSet a -> DisjointSet a
forall a b. a -> b -> b
`seq` Map a a -> Map a (RankChildren a) -> DisjointSet a
forall a. Map a a -> Map a (RankChildren a) -> DisjointSet a
DisjointSet Map a a
p' Map a (RankChildren a)
r

{- $tutorial

The disjoint set data structure represents sets that are
disjoint. Each set in the data structure can be interpreted
as an equivalance class. For example, let us consider a scenario
in which we are investigating spies who each use one or more aliases. There are three
actions we may repeated take:

    1. we learn an alias is in use by someone (make set)
    2. we learn two aliases refer to the same individual (union)
    3. we check our notes to figure out if two aliases refer to the same individual (find)

We initially learn of the existence of several aliases:

>>> import Data.Function ((&))
>>> import Data.Monoid ((<>))
>>> import Data.Foldable (fold,foldMap)
>>> let s0 = empty
>>> let s1 = s0 & insert "Scar" & insert "Carene" & insert "Barth" & insert "Coral"
>>> let s2 = s1 & insert "Boris" & insert "Esma" & insert "Mayra"
>>> putStr (pretty s2)
{{"Barth"},{"Boris"},{"Carene"},{"Coral"},{"Esma"},{"Mayra"},{"Scar"}}

Note that the 'Monoid' instance gives us a way to construct this more succintly:

>>> s2 == foldMap singleton ["Barth","Boris","Carene","Coral","Esma","Mayra","Scar"]
True

After some preliminary research, we learn that Barth and Esma are the same person. We
also learn that Carene and Mayra are the same:

>>> let s3 = s2 & union "Barth" "Esma" & union "Carene" "Mayra"
>>> putStr (pretty s3)
{{"Boris"},{"Coral"},{"Barth","Esma"},{"Carene","Mayra"},{"Scar"}}

Another informant comes forward who tells us they have worked for someone
that went by the names Mayra and Esma. Going through old letters, we learn
that Boris is a pen name used by Scar:

>>> let s4 = s3 & union "Mayra" "Esma" & union "Boris" "Scar"
>>> putStr (pretty s4)
{{"Coral"},{"Barth","Carene","Esma","Mayra"},{"Boris","Scar"}}

At this point, Detective Laura from another department drops by with
questions about a case she is working on. She asks if Boris the same
person as Barth and if Carene is the same person as Esma. We answer:

>>> equivalent "Boris" "Barth" s4
False
>>> equivalent "Carene" "Esma" s4
True

The correct way to interpret this is that @False@ means something more
along the lines of unknown, but we definitely know that Carene is Esma.
Finally, before the detective leaves, she gives us some of her case
notes to synthesize with our information. Notice that there are
some aliases she encountered that we did not and vice versa:

>>> let laura = union "Scar" "Coral" $ union "Esma" "Henri" $ foldMap singleton ["Carene","Boris","Barth"]
>>> putStr (pretty laura)
{{"Barth"},{"Boris"},{"Carene"},{"Coral","Scar"},{"Esma","Henri"}}
>>> putStr (pretty (laura <> s4))
{{"Barth","Carene","Esma","Henri","Mayra"},{"Boris","Coral","Scar"}}

With Laura's shared findings, we now see that there are really only (at most)
two spies that we are dealing with.
-}