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