@@ -7,102 +7,113 @@ public struct UnionFind<T: Hashable> {
77 private var size = [ Int] ( )
88
99
10- public mutating func addSetWithElement ( element: T ) {
10+ public mutating func addSetWith ( element: T ) {
1111 index [ element] = parent. count
1212 parent. append ( parent. count)
1313 size. append ( 1 )
1414 }
1515
16- private mutating func findSetByIndexOfElement ( index: Int ) -> Int {
16+ private mutating func setByIndex ( index: Int ) -> Int {
1717 if parent [ index] == index {
1818 return index
1919 } else {
20- parent [ index] = findSetByIndexOfElement ( parent [ index] )
20+ parent [ index] = setByIndex ( parent [ index] )
2121 return parent [ index]
2222 }
2323 }
2424
25- public mutating func findSetOfElement ( element: T ) -> Int ? {
25+ public mutating func setOf ( element: T ) -> Int ? {
2626 if let indexOfElement = index [ element] {
27- return findSetByIndexOfElement ( indexOfElement)
27+ return setByIndex ( indexOfElement)
2828 } else {
2929 return nil
3030 }
3131 }
3232
33- public mutating func unionSetsWithElement ( firstElement: T , andSecondElement secondElement: T ) {
34- if let firstSet = findSetOfElement ( firstElement) , secondSet = findSetOfElement ( secondElement) {
35- if ( firstSet != secondSet) {
36- if ( size [ firstSet] < size [ secondSet] ) {
37- parent [ firstSet] = secondSet;
33+ public mutating func unionSetsContaining ( firstElement: T , and secondElement: T ) {
34+ if let firstSet = setOf ( firstElement) , secondSet = setOf ( secondElement) {
35+ if firstSet != secondSet {
36+ if size [ firstSet] < size [ secondSet] {
37+ parent [ firstSet] = secondSet
3838 size [ secondSet] += size [ firstSet]
3939 } else {
40- parent [ secondSet] = firstSet;
40+ parent [ secondSet] = firstSet
4141 size [ firstSet] += size [ secondSet]
4242 }
4343 }
4444 }
4545 }
46+
47+ public mutating func inSameSet( firstElement: T, and secondElement: T) - > Bool {
48+ if let firstSet = setOf ( firstElement) , secondSet = setOf ( secondElement) {
49+ return firstSet == secondSet
50+ } else {
51+ return false
52+ }
53+ }
4654}
4755
4856
4957var dsu = UnionFind < Int > ( )
5058
5159for i in 1 ... 10 {
52- dsu. addSetWithElement ( i)
60+ dsu. addSetWith ( i)
5361}
5462// now our dsu contains 10 independent sets
5563
5664// let's divide our numbers into two sets by divisibility by 2
5765for i in 3 ... 10 {
5866 if i % 2 == 0 {
59- dsu. unionSetsWithElement ( 2 , andSecondElement : i)
67+ dsu. unionSetsContaining ( 2 , and : i)
6068 } else {
61- dsu. unionSetsWithElement ( 1 , andSecondElement : i)
69+ dsu. unionSetsContaining ( 1 , and : i)
6270 }
6371}
6472
6573// check our division
66- print( dsu. findSetOfElement ( 2 ) == dsu. findSetOfElement ( 4 ) )
67- print( dsu. findSetOfElement ( 4 ) == dsu. findSetOfElement ( 6 ) )
68- print( dsu. findSetOfElement ( 6 ) == dsu. findSetOfElement ( 8 ) )
69- print( dsu. findSetOfElement ( 8 ) == dsu. findSetOfElement ( 10 ) )
74+ print( dsu. inSameSet ( 2 , and: 4 ) )
75+ print( dsu. inSameSet ( 4 , and: 6 ) )
76+ print( dsu. inSameSet ( 6 , and: 8 ) )
77+ print( dsu. inSameSet ( 8 , and: 10 ) )
78+
79+
80+ print( dsu. inSameSet ( 1 , and: 3 ) )
81+ print( dsu. inSameSet ( 3 , and: 5 ) )
82+ print( dsu. inSameSet ( 5 , and: 7 ) )
83+ print( dsu. inSameSet ( 7 , and: 9 ) )
84+
7085
71- print( dsu. findSetOfElement ( 1 ) == dsu. findSetOfElement ( 3 ) )
72- print( dsu. findSetOfElement ( 3 ) == dsu. findSetOfElement ( 5 ) )
73- print( dsu. findSetOfElement ( 5 ) == dsu. findSetOfElement ( 7 ) )
74- print( dsu. findSetOfElement ( 7 ) == dsu. findSetOfElement ( 9 ) )
86+ print( dsu. inSameSet ( 7 , and: 4 ) )
87+ print( dsu. inSameSet ( 3 , and: 6 ) )
7588
76- print( dsu. findSetOfElement ( 8 ) == dsu. findSetOfElement ( 9 ) )
77- print( dsu. findSetOfElement ( 4 ) == dsu. findSetOfElement ( 3 ) )
7889
7990
8091var dsuForStrings = UnionFind < String > ( )
8192let words = [ " all " , " border " , " boy " , " afternoon " , " amazing " , " awesome " , " best " ]
8293
83- dsuForStrings. addSetWithElement ( " a " )
84- dsuForStrings. addSetWithElement ( " b " )
94+ dsuForStrings. addSetWith ( " a " )
95+ dsuForStrings. addSetWith ( " b " )
8596
8697// In that example we divide strings by its first letter
8798for word in words {
88- dsuForStrings. addSetWithElement ( word)
99+ dsuForStrings. addSetWith ( word)
89100 if word. hasPrefix ( " a " ) {
90- dsuForStrings. unionSetsWithElement ( " a " , andSecondElement : word)
101+ dsuForStrings. unionSetsContaining ( " a " , and : word)
91102 } else if word. hasPrefix ( " b " ) {
92- dsuForStrings. unionSetsWithElement ( " b " , andSecondElement : word)
103+ dsuForStrings. unionSetsContaining ( " b " , and : word)
93104 }
94105}
95106
96- print( dsuForStrings. findSetOfElement ( " a " ) == dsuForStrings . findSetOfElement ( " all " ) )
97- print( dsuForStrings. findSetOfElement ( " all " ) == dsuForStrings . findSetOfElement ( " awesome " ) )
98- print( dsuForStrings. findSetOfElement ( " amazing " ) == dsuForStrings . findSetOfElement ( " afternoon " ) )
107+ print( dsuForStrings. inSameSet ( " a " , and : " all " ) )
108+ print( dsuForStrings. inSameSet ( " all " , and : " awesome " ) )
109+ print( dsuForStrings. inSameSet ( " amazing " , and : " afternoon " ) )
99110
100- print( dsuForStrings. findSetOfElement ( " b " ) == dsuForStrings . findSetOfElement ( " boy " ) )
101- print( dsuForStrings. findSetOfElement ( " best " ) == dsuForStrings . findSetOfElement ( " boy " ) )
102- print( dsuForStrings. findSetOfElement ( " border " ) == dsuForStrings . findSetOfElement ( " best " ) )
111+ print( dsuForStrings. inSameSet ( " b " , and : " boy " ) )
112+ print( dsuForStrings. inSameSet ( " best " , and : " boy " ) )
113+ print( dsuForStrings. inSameSet ( " border " , and : " best " ) )
103114
104- print( dsuForStrings. findSetOfElement ( " amazing " ) == dsuForStrings . findSetOfElement ( " boy " ) )
105- print( dsuForStrings. findSetOfElement ( " all " ) == dsuForStrings . findSetOfElement ( " border " ) )
115+ print( dsuForStrings. inSameSet ( " amazing " , and : " boy " ) )
116+ print( dsuForStrings. inSameSet ( " all " , and : " border " ) )
106117
107118
108119
0 commit comments