blob: 3a8085f2946ae4d8b9c57877a0407bdaecdee1b1 [file] [log] [blame]
khenaidooac637102019-01-14 15:44:34 -05001package goraph
2
3import "sync"
4
5// DisjointSet implements disjoint set.
6// (https://en.wikipedia.org/wiki/Disjoint-set_data_structure)
7type DisjointSet struct {
8 represent string
9 members map[string]struct{}
10}
11
12// Forests is a set of DisjointSet.
13type Forests struct {
14 mu sync.Mutex // guards the following
15 data map[*DisjointSet]struct{}
16}
17
18// NewForests creates a new Forests.
19func NewForests() *Forests {
20 set := &Forests{}
21 set.data = make(map[*DisjointSet]struct{})
22 return set
23}
24
25// MakeDisjointSet creates a DisjointSet.
26func MakeDisjointSet(forests *Forests, name string) {
27 newDS := &DisjointSet{}
28 newDS.represent = name
29 members := make(map[string]struct{})
30 members[name] = struct{}{}
31 newDS.members = members
32 forests.mu.Lock()
33 defer forests.mu.Unlock()
34 forests.data[newDS] = struct{}{}
35}
36
37// FindSet returns the DisjointSet with the represent name.
38func FindSet(forests *Forests, name string) *DisjointSet {
39 forests.mu.Lock()
40 defer forests.mu.Unlock()
41 for data := range forests.data {
42 if data.represent == name {
43 return data
44 }
45 for k := range data.members {
46 if k == name {
47 return data
48 }
49 }
50 }
51 return nil
52}
53
54// Union unions two DisjointSet, with ds1's represent.
55func Union(forests *Forests, ds1, ds2 *DisjointSet) {
56 newDS := &DisjointSet{}
57 newDS.represent = ds1.represent
58 newDS.members = ds1.members
59 for k := range ds2.members {
60 newDS.members[k] = struct{}{}
61 }
62 forests.mu.Lock()
63 defer forests.mu.Unlock()
64 forests.data[newDS] = struct{}{}
65 delete(forests.data, ds1)
66 delete(forests.data, ds2)
67}