HeightedUnion.swift

struct HeightedUnion {
    
    typealias QuickUnionRootHeight = (root: Int, height: Int)
    
    private var ids: [Int]
    
    // Complexity: O(n)
    init(_ N: Int) {
        ids = Array(0..<N)
    }
    
    private func rootHeight(var i: Int) -> QuickUnionRootHeight {
        var height = 1
        while (i != ids[i]) {
            i = ids[i]
            height++
        }
        
        return (i, height)
    }
    
    // Complexity: O(n) - Worst carse we will traverse everything
    func connected(p: Int, q: Int) -> Bool {
        return  rootHeight(p).root == rootHeight(q).root
    }
    
    // Complexity: O(n) - Worst carse we will traverse everything
    mutating func union(from: Int, _ to: Int) {
        let rootHeightFrom = rootHeight(from)
        let rootHeightTo = rootHeight(to)
        
        if rootHeightFrom.height <= rootHeightTo.height {
            ids[from] = to
        }
    }
    
}