|
| 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 | + |
0 commit comments