Skip to content

Commit 743b2f6

Browse files
committed
Add Union-Find data structure implementation
1 parent b9a9759 commit 743b2f6

File tree

5 files changed

+170
-0
lines changed

5 files changed

+170
-0
lines changed

UnionFindDS/README.markdown

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
# Union Find Data Structure
2+
3+
TODO
Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
//: Playground - noun: a place where people can play
2+
3+
public class UnionFindDS<T: Hashable> {
4+
5+
private var index = [T:Int]()
6+
private var parent = [Int]()
7+
private var size = [Int]()
8+
9+
10+
public func addSetWithElement(element: T) {
11+
index[element] = parent.count
12+
parent.append(parent.count)
13+
size.append(1)
14+
}
15+
16+
private func findSetByIndexOfElement(index: Int) -> Int {
17+
if parent[index] == index {
18+
return index
19+
} else {
20+
parent[index] = findSetByIndexOfElement(parent[index])
21+
return parent[index]
22+
}
23+
}
24+
25+
public func findSetOfElement(element: T) -> Int {
26+
let indexOfElement = index[element]!
27+
return findSetByIndexOfElement(indexOfElement)
28+
}
29+
30+
public func unionSetsWithElement(firstElement: T, andSecondElement secondElement: T) {
31+
let firstSet = findSetOfElement(firstElement)
32+
let secondSet = findSetOfElement(secondElement)
33+
if (firstSet != secondSet) {
34+
if (size[firstSet] < size[secondSet]) {
35+
parent[firstSet] = secondSet;
36+
size[secondSet] += size[firstSet]
37+
} else {
38+
parent[secondSet] = firstSet;
39+
size[firstSet] += size[secondSet]
40+
}
41+
}
42+
}
43+
}
44+
45+
46+
let dsu = UnionFindDS<Int>()
47+
48+
for i in 1...10 {
49+
dsu.addSetWithElement(i)
50+
}
51+
// now our dsu contains 10 independent sets
52+
53+
// let's divide our numbers into two sets by divisibility by 2
54+
for i in 3...10 {
55+
if i % 2 == 0 {
56+
dsu.unionSetsWithElement(2, andSecondElement: i)
57+
} else {
58+
dsu.unionSetsWithElement(1, andSecondElement: i)
59+
}
60+
}
61+
62+
// check our division
63+
print(dsu.findSetOfElement(2) == dsu.findSetOfElement(4))
64+
print(dsu.findSetOfElement(4) == dsu.findSetOfElement(6))
65+
print(dsu.findSetOfElement(6) == dsu.findSetOfElement(8))
66+
print(dsu.findSetOfElement(8) == dsu.findSetOfElement(10))
67+
68+
print(dsu.findSetOfElement(1) == dsu.findSetOfElement(3))
69+
print(dsu.findSetOfElement(3) == dsu.findSetOfElement(5))
70+
print(dsu.findSetOfElement(5) == dsu.findSetOfElement(7))
71+
print(dsu.findSetOfElement(7) == dsu.findSetOfElement(9))
72+
73+
print(dsu.findSetOfElement(8) == dsu.findSetOfElement(9))
74+
print(dsu.findSetOfElement(4) == dsu.findSetOfElement(3))
75+
76+
77+
let dsuForStrings = UnionFindDS<String>()
78+
let words = ["all", "border", "boy", "afternoon", "amazing", "awesome", "best"]
79+
80+
dsuForStrings.addSetWithElement("a")
81+
dsuForStrings.addSetWithElement("b")
82+
83+
// In that example we divide strings by its first letter
84+
for word in words {
85+
dsuForStrings.addSetWithElement(word)
86+
if word.hasPrefix("a") {
87+
dsuForStrings.unionSetsWithElement("a", andSecondElement: word)
88+
} else if word.hasPrefix("b") {
89+
dsuForStrings.unionSetsWithElement("b", andSecondElement: word)
90+
}
91+
}
92+
93+
print(dsuForStrings.findSetOfElement("a") == dsuForStrings.findSetOfElement("all"))
94+
print(dsuForStrings.findSetOfElement("all") == dsuForStrings.findSetOfElement("awesome"))
95+
print(dsuForStrings.findSetOfElement("amazing") == dsuForStrings.findSetOfElement("afternoon"))
96+
97+
print(dsuForStrings.findSetOfElement("b") == dsuForStrings.findSetOfElement("boy"))
98+
print(dsuForStrings.findSetOfElement("best") == dsuForStrings.findSetOfElement("boy"))
99+
print(dsuForStrings.findSetOfElement("border") == dsuForStrings.findSetOfElement("best"))
100+
101+
print(dsuForStrings.findSetOfElement("amazing") == dsuForStrings.findSetOfElement("boy"))
102+
print(dsuForStrings.findSetOfElement("all") == dsuForStrings.findSetOfElement("border"))
103+
104+
105+
106+
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
2+
<playground version='5.0' target-platform='osx'>
3+
<timeline fileName='timeline.xctimeline'/>
4+
</playground>
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<Timeline
3+
version = "3.0">
4+
<TimelineItems>
5+
</TimelineItems>
6+
</Timeline>

UnionFindDS/UnionFindDS.swift

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
Union-Find Data Structure
3+
4+
Performance:
5+
adding new set is almost O(1)
6+
finding set of element is almost O(1)
7+
union sets is almost O(1)
8+
*/
9+
10+
11+
public class UnionFindDS<T: Hashable> {
12+
13+
private var index = [T:Int]()
14+
private var parent = [Int]()
15+
private var size = [Int]()
16+
17+
18+
public func addSetWithElement(element: T) {
19+
index[element] = parent.count
20+
parent.append(parent.count)
21+
size.append(1)
22+
}
23+
24+
private func findSetByIndexOfElement(index: Int) -> Int {
25+
if parent[index] == index {
26+
return index
27+
} else {
28+
parent[index] = findSetByIndexOfElement(parent[index])
29+
return parent[index]
30+
}
31+
}
32+
33+
public func findSetOfElement(element: T) -> Int {
34+
let indexOfElement = index[element]!
35+
return findSetByIndexOfElement(indexOfElement)
36+
}
37+
38+
public func unionSetsWithElement(firstElement: T, andSecondElement secondElement: T) {
39+
let firstSet = findSetOfElement(firstElement)
40+
let secondSet = findSetOfElement(secondElement)
41+
if (firstSet != secondSet) {
42+
if (size[firstSet] < size[secondSet]) {
43+
parent[firstSet] = secondSet;
44+
size[secondSet] += size[firstSet]
45+
} else {
46+
parent[secondSet] = firstSet;
47+
size[firstSet] += size[secondSet]
48+
}
49+
}
50+
}
51+
}

0 commit comments

Comments
 (0)