|
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@2219
|
6 val classes : unionFind -> K.ord_key list list
|
|
ziv@2219
|
7 end = struct
|
|
ziv@2216
|
8
|
|
ziv@2216
|
9 structure M = BinaryMapFn(K)
|
|
ziv@2216
|
10 structure S = BinarySetFn(K)
|
|
ziv@2216
|
11
|
|
ziv@2216
|
12 datatype entry =
|
|
ziv@2216
|
13 Set of S.set
|
|
ziv@2216
|
14 | Pointer of K.ord_key
|
|
ziv@2216
|
15
|
|
ziv@2216
|
16 (* First map is the union-find tree, second stores equivalence classes. *)
|
|
ziv@2216
|
17 type unionFind = entry M.map ref * S.set M.map
|
|
ziv@2216
|
18
|
|
ziv@2216
|
19 val empty : unionFind = (ref M.empty, M.empty)
|
|
ziv@2216
|
20
|
|
ziv@2216
|
21 fun findPair (uf, x) =
|
|
ziv@2216
|
22 case M.find (!uf, x) of
|
|
ziv@2216
|
23 NONE => (S.singleton x, x)
|
|
ziv@2216
|
24 | SOME (Set set) => (set, x)
|
|
ziv@2216
|
25 | SOME (Pointer parent) =>
|
|
ziv@2216
|
26 let
|
|
ziv@2216
|
27 val (set, rep) = findPair (uf, parent)
|
|
ziv@2216
|
28 in
|
|
ziv@2216
|
29 uf := M.insert (!uf, x, Pointer rep);
|
|
ziv@2216
|
30 (set, rep)
|
|
ziv@2216
|
31 end
|
|
ziv@2216
|
32
|
|
ziv@2216
|
33 fun find ((uf, _), x) = (S.listItems o #1 o findPair) (uf, x)
|
|
ziv@2216
|
34
|
|
ziv@2216
|
35 fun classes (_, cs) = (map S.listItems o M.listItems) cs
|
|
ziv@2216
|
36
|
|
ziv@2216
|
37 fun union ((uf, cs), x, y) =
|
|
ziv@2216
|
38 let
|
|
ziv@2216
|
39 val (xSet, xRep) = findPair (uf, x)
|
|
ziv@2216
|
40 val (ySet, yRep) = findPair (uf, y)
|
|
ziv@2216
|
41 val xySet = S.union (xSet, ySet)
|
|
ziv@2216
|
42 in
|
|
ziv@2216
|
43 (ref (M.insert (M.insert (!uf, yRep, Pointer xRep),
|
|
ziv@2216
|
44 xRep, Set xySet)),
|
|
ziv@2216
|
45 M.insert (case M.find (cs, yRep) of
|
|
ziv@2216
|
46 NONE => cs
|
|
ziv@2216
|
47 | SOME _ => #1 (M.remove (cs, yRep)),
|
|
ziv@2216
|
48 xRep, xySet))
|
|
ziv@2216
|
49 end
|
|
ziv@2216
|
50
|
|
ziv@2216
|
51 fun union' ((x, y), uf) = union (uf, x, y)
|
|
ziv@2216
|
52
|
|
ziv@2216
|
53 end
|
