union-find 0.1 → 0.2
raw patch · 5 files changed
+359/−4 lines, 5 filesdep +containersdep +transformersdep ~base
Dependencies added: containers, transformers
Dependency ranges changed: base
Files
- README.md +25/−2
- src/Control/Monad/Trans/UnionFind.hs +59/−0
- src/Data/UnionFind/IntMap.hs +101/−0
- src/Data/UnionFind/ST.hs +169/−0
- union-find.cabal +5/−2
README.md view
@@ -1,2 +1,25 @@-union-find-========+# union-find++A simple Haskell library that implements Tarjan's Union/Find+algorithm. Useful, for example, to implement unification in a type+inference system.++The Union/Find algorithm implements these operations in+(effectively) constant-time:++ 1. Check whether two elements are in the same equivalence class.++ 2. Create a union of two equivalence classes.++ 3. Look up the descriptor of the equivalence class.+++## Installation++Using `cabal` (which comes with the Haskell Platform):++ $ cabal install union-find++or in the checked-out repository:++ $ cabal install
+ src/Control/Monad/Trans/UnionFind.hs view
@@ -0,0 +1,59 @@+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+module Control.Monad.Trans.UnionFind+ ( UnionFindT, runUnionFind+ , Point, fresh, repr, descriptor, union, equivalent+ ) where++import Control.Applicative (Applicative)+import Control.Monad.Trans.Class (MonadTrans(..))+import Control.Monad.Trans.State (StateT(..), evalStateT)+import Data.UnionFind.IntMap (Point)+import qualified Control.Monad.Trans.State as State+import qualified Data.UnionFind.IntMap as UF++-- | A monad transformer that adds union find operations.+--+-- The @p@ parameter is the type of points. Uses the+-- "Data.UnionFind.IntMap" as the underlying union-find+-- implementation.+newtype UnionFindT p m a = UnionFindT {+ unUnionFindT :: StateT (UF.PointSupply p) m a+ } deriving (Functor, Applicative, Monad, MonadTrans)++runUnionFind :: Monad m => UnionFindT p m a -> m a+runUnionFind = (`evalStateT` UF.newPointSupply) . unUnionFindT++swap :: (a, b) -> (b, a)+swap (x, y) = (y, x)++-- | Create a new point with the given descriptor. The returned is+-- only equivalent to itself.+--+-- Note that a 'Point' has its own identity. That is, if two points+-- are equivalent then their descriptors are equal, but not vice+-- versa.+--+fresh :: Monad m => p -> UnionFindT p m (Point p)+fresh x = UnionFindT . StateT $ return . swap . flip UF.fresh x++-- | /O(1)/. @repr point@ returns the representative point of+-- @point@'s equivalence class.+repr :: Monad m => Point p -> UnionFindT p m (Point p)+repr = UnionFindT . State.gets . flip UF.repr++-- | Return the descriptor of the +descriptor :: Monad m => Point p -> UnionFindT p m p+descriptor = UnionFindT . State.gets . flip UF.descriptor++-- | Join the equivalence classes of the points. The resulting+-- equivalence class will get the descriptor of the second argument.+union :: Monad m => Point p -> Point p -> UnionFindT p m ()+union p1 p2 = UnionFindT . State.modify $ \x -> UF.union x p1 p2++-- | Test if the two elements are in the same equivalence class.+-- +-- @+-- liftA2 (==) (repr x) (repr y)+-- @+equivalent :: Monad m => Point p -> Point p -> UnionFindT p m Bool+equivalent p1 p2 = UnionFindT . State.gets $ \x -> UF.equivalent x p1 p2
+ src/Data/UnionFind/IntMap.hs view
@@ -0,0 +1,101 @@+{-# LANGUAGE Rank2Types #-}+{-# LANGUAGE BangPatterns #-}+{-# OPTIONS_GHC -funbox-strict-fields #-}+--+-- NOTE: Does not currently ensure that a 'Point' was indeed generated+-- by the specified 'PointSupply'.+--+module Data.UnionFind.IntMap + ( newPointSupply, fresh, repr, descriptor, union, equivalent,+ PointSupply, Point ) where++import qualified Data.IntMap as IM++data PointSupply a = PointSupply !Int (IM.IntMap (Link a))+ deriving Show++data Link a + = Info {-# UNPACK #-} !Int a+ -- ^ This is the descriptive element of the equivalence class+ -- and its rank.+ | Link {-# UNPACK #-} !Int+ -- ^ Pointer to some other element of the equivalence class.+ deriving Show++newtype Point a = Point Int++newPointSupply :: PointSupply a+newPointSupply = PointSupply 0 IM.empty++fresh :: PointSupply a -> a -> (PointSupply a, Point a)+fresh (PointSupply next eqs) a =+ (PointSupply (next + 1) (IM.insert next (Info 0 a) eqs), Point next)++-- freshList :: PointSupply a -> [a] -> (PointSupply a, [Point a])+-- freshList ++repr :: PointSupply a -> Point a -> Point a+repr ps p = reprInfo ps p (\n _rank _a -> Point n)++reprInfo :: PointSupply a -> Point a -> (Int -> Int -> a -> r) -> r+reprInfo (PointSupply _next eqs) (Point n) k = go n+ where+ go !i =+ case eqs IM.! i of+ Link i' -> go i'+ Info r a -> k i r a+ +union :: PointSupply a -> Point a -> Point a -> PointSupply a+union ps@(PointSupply next eqs) p1 p2 =+ reprInfo ps p1 $ \i1 r1 _a1 -> + reprInfo ps p2 $ \i2 r2 a2 ->+ if i1 == i2 then ps else+ case r1 `compare` r2 of+ LT ->+ -- No rank or descriptor update necessary+ let !eqs1 = IM.insert i1 (Link i2) eqs in+ PointSupply next eqs1+ EQ ->+ let !eqs1 = IM.insert i1 (Link i2) eqs+ !eqs2 = IM.insert i2 (Info (r2 + 1) a2) eqs1 in+ PointSupply next eqs2+ GT ->+ let !eqs1 = IM.insert i1 (Info r2 a2) eqs+ !eqs2 = IM.insert i2 (Link i1) eqs1 in+ PointSupply next eqs2++descriptor :: PointSupply a -> Point a -> a+descriptor ps p = reprInfo ps p (\_ _ a -> a)++equivalent :: PointSupply a -> Point a -> Point a -> Bool+equivalent ps p1 p2 =+ reprInfo ps p1 $ \i1 _ _ ->+ reprInfo ps p2 $ \i2 _ _ ->+ i1 == i2++{-+tst1 :: IO ()+tst1 = do+ let ps0 = newPointSupply+ (ps1, p1) = fresh ps0 "hello"+ (ps2, p2) = fresh ps1 "world"+ (ps3, p3) = fresh ps2 "you"+ (ps, p4) = fresh ps3 "there"+ let ps' = union ps p1 p2+ print (descr ps p1, descr ps p2, equivalent ps p1 p2)+ print (descr ps' p1, descr ps' p2, equivalent ps' p1 p2)+ let ps'' = union ps' p3 p1+ print (descr ps'' p1, descr ps'' p3, equivalent ps'' p1 p3)+ print ps''+-}++-- TODO: should fail+{-+tst2 :: IO ()+tst2 = do+ let as0 = newPointSupply+ (as, a1) = fresh as0 "foo"+ bs0 = newPointSupply+ (bs, b1) = fresh bs0 "bar"+ print $ union as a1 b1+ -}
+ src/Data/UnionFind/ST.hs view
@@ -0,0 +1,169 @@+-- | An implementation of Tarjan's UNION-FIND algorithm. (Robert E+-- Tarjan. \"Efficiency of a Good But Not Linear Set Union Algorithm\", JACM+-- 22(2), 1975)+--+-- The algorithm implements three operations efficiently (all amortised+-- @O(1)@):+--+-- 1. Check whether two elements are in the same equivalence class.+--+-- 2. Create a union of two equivalence classes.+--+-- 3. Look up the descriptor of the equivalence class.+-- +-- The implementation is based on mutable references. Each+-- equivalence class has exactly one member that serves as its+-- representative element. Every element either is the representative+-- element of its equivalence class or points to another element in+-- the same equivalence class. Equivalence testing thus consists of+-- following the pointers to the representative elements and then+-- comparing these for identity.+--+-- The algorithm performs lazy path compression. That is, whenever we+-- walk along a path greater than length 1 we automatically update the+-- pointers along the path to directly point to the representative+-- element. Consequently future lookups will be have a path length of+-- at most 1.+--+{-# OPTIONS_GHC -funbox-strict-fields #-}+module Data.UnionFind.ST+ ( Point, fresh, repr, union, union', equivalent, redundant,+ descriptor, setDescriptor, modifyDescriptor )+where++import Control.Applicative+import Control.Monad ( when )+import Control.Monad.ST+import Data.STRef++-- | The abstract type of an element of the sets we work on. It is+-- parameterised over the type of the descriptor.+newtype Point s a = Pt (STRef s (Link s a)) deriving Eq++data Link s a+ = Info {-# UNPACK #-} !(STRef s (Info a))+ -- ^ This is the descriptive element of the equivalence class.+ | Link {-# UNPACK #-} !(Point s a)+ -- ^ Pointer to some other element of the equivalence class.+ deriving Eq++data Info a = MkInfo+ { weight :: {-# UNPACK #-} !Int+ -- ^ The size of the equivalence class, used by 'union'.+ , descr :: a+ } deriving Eq++-- | /O(1)/. Create a fresh point and return it. A fresh point is in+-- the equivalence class that contains only itself.+fresh :: a -> ST s (Point s a)+fresh desc = do+ info <- newSTRef (MkInfo { weight = 1, descr = desc })+ l <- newSTRef (Info info)+ return (Pt l)++-- | /O(1)/. @repr point@ returns the representative point of+-- @point@'s equivalence class.+--+-- This method performs the path compresssion.+repr :: Point s a -> ST s (Point s a)+repr point@(Pt l) = do+ link <- readSTRef l+ case link of+ Info _ -> return point+ Link pt'@(Pt l') -> do+ pt'' <- repr pt'+ when (pt'' /= pt') $ do+ -- At this point we know that @pt'@ is not the representative+ -- element of @point@'s equivalent class. Therefore @pt'@'s+ -- link must be of the form @Link r@. We write this same+ -- value into @point@'s link reference and thereby perform+ -- path compression.+ link' <- readSTRef l'+ writeSTRef l link'+ return pt''++-- | Return the reference to the point's equivalence class's+-- descriptor.+descrRef :: Point s a -> ST s (STRef s (Info a))+descrRef point@(Pt link_ref) = do+ link <- readSTRef link_ref+ case link of+ Info info -> return info+ Link (Pt link'_ref) -> do+ link' <- readSTRef link'_ref+ case link' of+ Info info -> return info+ _ -> descrRef =<< repr point++-- | /O(1)/. Return the descriptor associated with argument point's+-- equivalence class.+descriptor :: Point s a -> ST s a+descriptor point = do+ descr <$> (readSTRef =<< descrRef point)++-- | /O(1)/. Replace the descriptor of the point's equivalence class+-- with the second argument.+setDescriptor :: Point s a -> a -> ST s ()+setDescriptor point new_descr = do+ r <- descrRef point+ modifySTRef r $ \i -> i { descr = new_descr }++modifyDescriptor :: Point s a -> (a -> a) -> ST s ()+modifyDescriptor point f = do+ r <- descrRef point+ modifySTRef r $ \i -> i { descr = f (descr i) }++-- | /O(1)/. Join the equivalence classes of the points (which must be+-- distinct). The resulting equivalence class will get the descriptor+-- of the second argument.+union :: Point s a -> Point s a -> ST s ()+union p1 p2 = union' p1 p2 (\_ d2 -> return d2)++-- | Like 'union', but sets the descriptor returned from the callback.+-- +-- The intention is to keep the descriptor of the second argument to+-- the callback, but the callback might adjust the information of the+-- descriptor or perform side effects.+union' :: Point s a -> Point s a -> (a -> a -> ST s a) -> ST s ()+union' p1 p2 update = do+ point1@(Pt link_ref1) <- repr p1+ point2@(Pt link_ref2) <- repr p2+ -- The precondition ensures that we don't create cyclic structures.+ when (point1 /= point2) $ do+ Info info_ref1 <- readSTRef link_ref1+ Info info_ref2 <- readSTRef link_ref2+ MkInfo w1 d1 <- readSTRef info_ref1 -- d1 is discarded+ MkInfo w2 d2 <- readSTRef info_ref2+ d2' <- update d1 d2+ -- Make the smaller tree a a subtree of the bigger one. The idea+ -- is this: We increase the path length of one set by one.+ -- Assuming all elements are accessed equally often, this means+ -- the penalty is smaller if we do it for the smaller set of the+ -- two.+ if w1 >= w2 then do+ writeSTRef link_ref2 (Link point1)+ writeSTRef info_ref1 (MkInfo (w1 + w2) d2')+ else do+ writeSTRef link_ref1 (Link point2)+ writeSTRef info_ref2 (MkInfo (w1 + w2) d2')++-- | /O(1)/. Return @True@ if both points belong to the same+-- | equivalence class.+equivalent :: Point s a -> Point s a -> ST s Bool+equivalent p1 p2 = (==) <$> repr p1 <*> repr p2++-- | /O(1)/. Returns @True@ for all but one element of an equivalence+-- class. That is, if @ps = [p1, .., pn]@ are all in the same+-- equivalence class, then the following assertion holds.+-- +-- > do rs <- mapM redundant ps+-- > assert (length (filter (==False) rs) == 1)+-- +-- It is unspecified for which element function returns @False@, so be+-- really careful when using this.+redundant :: Point s a -> ST s Bool+redundant (Pt link_r) = do+ link <- readSTRef link_r+ case link of+ Info _ -> return False+ Link _ -> return True
union-find.cabal view
@@ -1,5 +1,5 @@ Name: union-find-Version: 0.1+Version: 0.2 License: BSD3 License-File: LICENSE Author: Thomas Schilling <nominolo@googlemail.com>@@ -28,10 +28,13 @@ Library Build-Depends:- base >= 4 && < 5+ base >= 4 && < 5, containers >= 0.3, transformers >= 0.2 GHC-Options: -Wall Exposed-Modules:+ Control.Monad.Trans.UnionFind Data.UnionFind.IO+ Data.UnionFind.ST+ Data.UnionFind.IntMap Hs-Source-Dirs: src