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
|