Add UnionFind data structure and tests

evgeniyd · Feb 26, 2024 · 0586796 · 0586796
1 parent 4a7f5e3
commit 0586796
Show file tree

Hide file tree

Showing 3 changed files with 157 additions and 4 deletions.
diff --git a/UnionStreet.xcodeproj/project.pbxproj b/UnionStreet.xcodeproj/project.pbxproj
@@ -31,6 +31,7 @@
 		584330592B8D21900072AFEC /* UnionStreetTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = UnionStreetTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
 		5843305E2B8D21900072AFEC /* UnionStreetTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UnionStreetTests.swift; sourceTree = "<group>"; };
 		5843306B2B8D22040072AFEC /* UnionFind.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UnionFind.swift; sourceTree = "<group>"; };
+		5843306D2B8D23CD0072AFEC /* README.md */ = {isa = PBXFileReference; lastKnownFileType = net.daringfireball.markdown; path = README.md; sourceTree = "<group>"; };
 /* End PBXFileReference section */
 
 /* Begin PBXFrameworksBuildPhase section */
@@ -55,6 +56,7 @@
 		584330452B8D218F0072AFEC = {
 			isa = PBXGroup;
 			children = (
+				5843306D2B8D23CD0072AFEC /* README.md */,
 				584330512B8D218F0072AFEC /* UnionStreet */,
 				5843305D2B8D21900072AFEC /* UnionStreetTests */,
 				584330502B8D218F0072AFEC /* Products */,

diff --git a/UnionStreet/UnionFind.swift b/UnionStreet/UnionFind.swift
@@ -8,5 +8,39 @@
 import Foundation
 
 struct UnionFind {
+    private var par: [Int]
+    private var rank: [Int]
 
+    init(_ count: Int) {
+        par = Array(0..<count)
+        rank = Array(repeating:1, count: count)
+    }
+
+    func find(_ n: Int) -> Int {
+        var n = par[n]
+        while n != par[n] {
+            n = par[n]
+        }
+        return n
+    }
+
+    @discardableResult
+    mutating func union(_ n1: Int, _ n2: Int) -> Bool {
+        let p1 = find(n1)
+        let p2 = find(n2)
+
+        if p1 != p2 {
+            if rank[p1] >= rank[p2] {
+                par[p2] = p1
+                rank[p1] += 1
+            } else {
+                par[p1] = p2
+                rank[p2] += 1
+            }
+
+            return true
+        }
+
+        return false
+    }
 }
diff --git a/UnionStreetTests/UnionStreetTests.swift b/UnionStreetTests/UnionStreetTests.swift
@@ -11,13 +11,130 @@ import XCTest
 final class UnionStreetTests: XCTestCase {
 
     func test_init_returnsNotNil() {
-        let uf = makeSUT()
-        XCTAssertNotNil(uf)
+        let sut = makeSUT()
+        XCTAssertNotNil(sut)
+    }
+
+    func test_find_singleNodeReturnsNodeItselfOnInit() {
+        let sut = makeSUT(1)
+        XCTAssertEqual(0, sut.find(0))
+    }
+
+    func test_union_returnsFalseOnUnionWithItself() {
+        var sut = makeSUT(1)
+        XCTAssertEqual(false, sut.union(0,0))
+    }
+
+    func test_find_returnsNodesItselfOnTwoNodesInit() {
+        let sut = makeSUT(2)
+        XCTAssertEqual(0, sut.find(0))
+        XCTAssertEqual(1, sut.find(1))
+    }
+
+    func test_union_returnsTrueOnUnionOfTwoNodesOnce() {
+        var sut = makeSUT(2)
+        XCTAssertEqual(true, sut.union(0,1))
+    }
+
+    func test_union_returnsFalseOnUnionOfTwoNodesTwice() {
+        var sut = makeSUT(2)
+
+        sut.union(0,1)
+
+        XCTAssertEqual(false, sut.union(0,1))
+    }
+
+    func test_union_returnsTrueOnUnionOfMultipleNodes() {
+        let input = [0,1,2,3,4,5]
+        var sut = makeSUT(input.count)
+        let parent = input.first!
+
+        for num in input where num != parent {
+            XCTAssertEqual(true, sut.union(parent,num))
+        }
+    }
+
+    func test_find_returnsSameParentOnMultipleNodesUnion() {
+        let input = [0,1,2,3,4,5]
+        var sut = makeSUT(input.count)
+        let parent = input.first!
+
+        for num in input {
+            sut.union(parent,num)
+        }
+
+        for num in input {
+            let receivedParent = sut.find(num)
+            let expectedParent = parent
+            XCTAssertEqual(receivedParent, expectedParent)
+        }
+    }
+
+    func test_union_returnsTrueOnUnionOfMultipleNodesWithArbitraryUniqueValues() {
+        let input = [6,3,4,1,5,0]
+        var sut = makeSUT(input.max()!+1)
+        let parent = input.first!
+
+        var expectedReturnValues = [false, true, true, true, true, true]
+
+        var receivedReturnValues: [Bool] = []
+        for num in input {
+            receivedReturnValues.append(sut.union(parent,num))
+        }
+
+        XCTAssertEqual(expectedReturnValues, receivedReturnValues, "Expecting \(expectedReturnValues), received \(receivedReturnValues)  instead")
+    }
+
+    func test_union_returnsTrueOnUnionOfMultipleNodesWithArbitraryValuesWithDuplicates() {
+        let input = [6,3,1,1,5,0]
+        var sut = makeSUT(input.max()!+1)
+        let parent = input.first!
+
+        var expectedReturnValues = [false, true, true, false, true, true]
+
+        var receivedReturnValues: [Bool] = []
+        for num in input {
+            receivedReturnValues.append(sut.union(parent,num))
+        }
+
+        XCTAssertEqual(expectedReturnValues, receivedReturnValues, "Expecting \(expectedReturnValues), received \(receivedReturnValues)  instead")
+    }
+
+    func test_find_returnsSameParentOnUnionOfMultipleNodesWithArbitraryUniqueValues() {
+        let input = [6,3,4,1,5,0]
+        var sut = makeSUT(input.max()!+1)
+        let parent = input.first!
+
+        for num in input {
+            sut.union(parent,num)
+        }
+
+        for num in input {
+            let receivedParent = sut.find(num)
+            let expectedParent = parent
+            XCTAssertEqual(receivedParent, expectedParent)
+        }
+    }
+
+    func test_find_returnsSameParentOnUnionOfMultipleNodesWithArbitraryValuesWithDuplicates() {
+        let input = [6,3,1,1,5,0]
+        var sut = makeSUT(input.max()!+1)
+        let parent = input.first!
+
+        for num in input {
+            sut.union(parent,num)
+        }
+
+        for num in input {
+            let receivedParent = sut.find(num)
+            let expectedParent = parent
+            XCTAssertEqual(receivedParent, expectedParent)
+        }
     }
 
     // MARK: Helpers
 
-    private func makeSUT() -> UnionFind {
-        return UnionFind()
+    private func makeSUT(_ count: Int = 0) -> UnionFind {
+        return UnionFind(count)
     }
 }