{-# LANGUAGE BangPatterns #-}

{- |
Module      : Data.IntDisjointSet
Description : Persistent Disjoint-Sets (a.k.a. Union-Find)
Copyright   : (c) 2012 Maxwell Sayles.
License     : LGPL

Maintainer  : maxwellsayles@gmail.com
Stability   : stable
Portability : non-portable (only tested with GHC 6.12.3)

Persistent Disjoint-Sets (a.k.a. Union-Find).  This implements
disjoint-sets according to the description given in
\"/Introduction to Algorithms/\" by Cormen et al
(<http://mitpress.mit.edu/algorithms>).
Most functions incur an additional O(logn) overhead due to the use
of persistent maps.

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. Each element has a set representative. The implementation works by
maintaining a map from an element to its parent. When an element is its
own parent, it is the set representative. Two elements are part of the
same equivalence set when their set representatives are the same.

In order to find the set representative efficiently, after each traversal
from an element to its representative, we compress the path so that
each element on the path points directly to the set representative.  For
this to be persistent, lookup is stateful and so returns the result
of the lookup and a new disjoint set.

Additionally, to make sure that path lengths grow logarithmically, we
maintain the rank of a set. This is a logarithmic upper bound on the 
number of elements in each set. When we compute the union of two sets,
we make the set with the smaller rank a child of the set with the larger
rank. When two sets have equal rank, the first set is a child of the second
and the rank of the second is increased by 1.
-}

module Data.IntDisjointSet (IntDisjointSet,
                            empty,
                            singleton,
                            insert,
                            unsafeMerge,
                            union,
                            lookup,
                            elems,
                            toList,
                            fromList,
                            equivalent,
                            disjointSetSize,
                            size,
                            map) where

import Control.Arrow
import Control.Monad.State.Strict
import Control.Monad.Trans.Maybe
import qualified Data.IntMap as IntMap
import qualified Data.List as List
import Data.Maybe
import Prelude hiding (lookup, map)

{-|
Represents a disjoint set of integers.
-}
data IntDisjointSet = IntDisjointSet { parents :: IntMap.IntMap Int,
                                       ranks   :: IntMap.IntMap Int }

instance Show IntDisjointSet where
    show = ("fromList " ++) . show . fst . toList

{-|
Create a disjoint set with no members.
O(1).
-}
empty :: IntDisjointSet
empty = IntDisjointSet IntMap.empty IntMap.empty

{-|
Create a disjoint set with one member.
O(1).
-}
singleton :: Int -> IntDisjointSet
singleton !x = let p = IntMap.singleton x x
                   r = IntMap.singleton x 0
               in  p `seq` r `seq` IntDisjointSet p r

{-|
Insert x into the disjoint set.
If it is already a member, then do nothing,
otherwise x has no equivalence relations.
O(logn).
-}
insert :: Int -> IntDisjointSet -> IntDisjointSet
insert !x set@(IntDisjointSet p r) =
    let (l, p') = IntMap.insertLookupWithKey (\_ _ old -> old) x x p
    in  case l of
          Just _  -> set
          Nothing ->
              let r' = IntMap.insert x 0 r
              in  p' `seq` r' `seq` IntDisjointSet p' r'
                    
{-|
Given two instances of disjoint sets that share no members in common,
computes a third disjoint set that is the combination of the two.

This method is unsafe in that is does not verify that the two input
sets share no members in common and in the event that a member
overlaps, the resulting set may have incorrect equivalence relations.
-}
unsafeMerge :: IntDisjointSet -> IntDisjointSet -> IntDisjointSet
unsafeMerge (IntDisjointSet p1 r1) (IntDisjointSet p2 r2) =
    IntDisjointSet (IntMap.union p1 p2) (IntMap.union r1 r2)

{-|
Create an equivalence relation between x and y.

Amortized O(logn * \alpha(n))
where \alpha(n) is the extremely slowly growing
inverse Ackermann function.

This function works by looking up the set representatives
for both x and y.  If they are the same, it does nothing.
Then it looks up the rank for both representatives and
makes the tree of the smaller ranked representative a
child of the tree of the larger ranked representative.
If both representatives have the same rank, x is made a
child of y and the rank of y is increase by 1.
-}
union :: Int -> Int -> IntDisjointSet -> IntDisjointSet
union !x !y set = flip execState set $ runMaybeT $ do
  repx <- MaybeT $ state $ lookup x
  repy <- MaybeT $ state $ lookup y
  guard $ repx /= repy
  (IntDisjointSet p r) <- get
  let rankx = r IntMap.! repx
  let ranky = r IntMap.! repy
  put $! case compare rankx ranky of
    LT -> let p' = IntMap.insert repx repy p
              r' = IntMap.delete repx r
          in  p' `seq` r' `seq` IntDisjointSet p' r'
    GT -> let p' = IntMap.insert repy repx p
              r' = IntMap.delete repy r
          in  p' `seq` r' `seq` IntDisjointSet p' r'
    EQ -> let p' = IntMap.insert repx repy p
              r' = IntMap.delete repx $! IntMap.insert repy (ranky + 1) r
          in  p' `seq` r' `seq` IntDisjointSet p' r'

{-|
Find the set representative for this input.
This performs path compression and so is stateful.

Amortized O(logn * \alpha(n))
where \alpha(n) is the extremely slowly growing
inverse Ackermann function.
-}
lookup :: Int -> IntDisjointSet -> (Maybe Int, IntDisjointSet)
lookup !x set =
  case find x set of
    Nothing  -> (Nothing, set)
    Just rep -> let set' = compress x rep set
                in  set' `seq` (Just rep, set')

{-| Return a list of all the elements. -}
-- This is stateful for consistency and possible future revisions.
elems :: IntDisjointSet -> ([Int], IntDisjointSet)
elems = IntMap.keys . parents &&& id

{-| Generate an association list of each element and its representative. -}
toList :: IntDisjointSet -> ([(Int, Int)], IntDisjointSet)
toList set = flip runState set $ do
               xs <- state elems
               forM xs $ \x -> do
                 Just rep <- state $ lookup x
                 return (x, rep)

{-|
Given an association list representing equivalences between elements,
generate the corresponding disjoint-set.
-}
fromList :: [(Int, Int)] -> IntDisjointSet
fromList = foldr (\(x, y) -> union x y . insert y . insert x) empty

{-| True if both elements belong to the same set. -}
equivalent :: Int -> Int -> IntDisjointSet -> (Bool, IntDisjointSet)
equivalent !x !y set = first (fromMaybe False) $
                       flip runState set $
                       runMaybeT $ do
                         repx <- MaybeT $ state $ lookup x
                         repy <- MaybeT $ state $ lookup y
                         return $! repx == repy

{-|
Return the number of disjoint sets.
O(1).
-}
disjointSetSize :: IntDisjointSet -> Int
disjointSetSize = IntMap.size . ranks

{-|
Return the number of elements in all disjoint sets.
O(1).
-}
size :: IntDisjointSet -> Int
size = IntMap.size . parents

{-|
Map each member to another Int.
The map function must be a bijection, i.e. 1-to-1 mapping.
-}
map :: (Int -> Int) -> IntDisjointSet -> IntDisjointSet
map f (IntDisjointSet p r) =
  let p' = IntMap.fromList $ List.map (f *** f) $ IntMap.toList p
      r' = IntMap.fromList $ List.map (first f) $ IntMap.toList r
  in  p' `seq` r' `seq` IntDisjointSet p' r'

-- Find the set representative.
-- This traverses parents until the parent of y == y and returns y.
find :: Int -> IntDisjointSet -> Maybe Int
find !x (IntDisjointSet p _) =
  do x' <- IntMap.lookup x p
     return $! if x == x' then x' else find' x'
  where find' y = let y' = p IntMap.! y
                  in  if y == y' then y' else find' y'

-- Given a start node and its representative, compress
-- the path to the root.
compress :: Int -> Int -> IntDisjointSet -> IntDisjointSet
compress !x !rep set = helper x set
    where helper !x set@(IntDisjointSet p r)
              | x == rep  = set
              | otherwise = helper x' set'
              where x'    = p IntMap.! x
                    set'  = let p' = IntMap.insert x rep p
                            in  p' `seq` IntDisjointSet p' r