annotate src/union_find_fn.sml @ 2278:b7615e0ac4b0

Fix bug in and clean up free path code.
author Ziv Scully <ziv@mit.edu>
date Tue, 10 Nov 2015 12:35:00 -0500
parents 0730e217fc9c
children
rev   line source
ziv@2219 1 functor UnionFindFn(K : ORD_KEY) :> sig
ziv@2219 2 type unionFind
ziv@2219 3 val empty : unionFind
ziv@2219 4 val union : unionFind * K.ord_key * K.ord_key -> unionFind
ziv@2219 5 val union' : (K.ord_key * K.ord_key) * unionFind -> unionFind
ziv@2274 6 val together : unionFind * K.ord_key * K.ord_key -> bool
ziv@2219 7 val classes : unionFind -> K.ord_key list list
ziv@2219 8 end = struct
ziv@2216 9
ziv@2216 10 structure M = BinaryMapFn(K)
ziv@2216 11 structure S = BinarySetFn(K)
ziv@2216 12
ziv@2216 13 datatype entry =
ziv@2216 14 Set of S.set
ziv@2216 15 | Pointer of K.ord_key
ziv@2216 16
ziv@2216 17 (* First map is the union-find tree, second stores equivalence classes. *)
ziv@2216 18 type unionFind = entry M.map ref * S.set M.map
ziv@2216 19
ziv@2216 20 val empty : unionFind = (ref M.empty, M.empty)
ziv@2216 21
ziv@2216 22 fun findPair (uf, x) =
ziv@2216 23 case M.find (!uf, x) of
ziv@2216 24 NONE => (S.singleton x, x)
ziv@2216 25 | SOME (Set set) => (set, x)
ziv@2216 26 | SOME (Pointer parent) =>
ziv@2216 27 let
ziv@2216 28 val (set, rep) = findPair (uf, parent)
ziv@2216 29 in
ziv@2216 30 uf := M.insert (!uf, x, Pointer rep);
ziv@2216 31 (set, rep)
ziv@2216 32 end
ziv@2216 33
ziv@2216 34 fun find ((uf, _), x) = (S.listItems o #1 o findPair) (uf, x)
ziv@2216 35
ziv@2216 36 fun classes (_, cs) = (map S.listItems o M.listItems) cs
ziv@2216 37
ziv@2274 38 fun together ((uf, _), x, y) = case K.compare (#2 (findPair (uf, x)), #2 (findPair (uf, y))) of
ziv@2274 39 EQUAL => true
ziv@2274 40 | _ => false
ziv@2274 41
ziv@2216 42 fun union ((uf, cs), x, y) =
ziv@2216 43 let
ziv@2216 44 val (xSet, xRep) = findPair (uf, x)
ziv@2216 45 val (ySet, yRep) = findPair (uf, y)
ziv@2216 46 val xySet = S.union (xSet, ySet)
ziv@2216 47 in
ziv@2216 48 (ref (M.insert (M.insert (!uf, yRep, Pointer xRep),
ziv@2216 49 xRep, Set xySet)),
ziv@2216 50 M.insert (case M.find (cs, yRep) of
ziv@2216 51 NONE => cs
ziv@2216 52 | SOME _ => #1 (M.remove (cs, yRep)),
ziv@2216 53 xRep, xySet))
ziv@2216 54 end
ziv@2216 55
ziv@2216 56 fun union' ((x, y), uf) = union (uf, x, y)
ziv@2216 57
ziv@2216 58 end