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 !Int a
| Link !Int
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)
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 ->
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