disjoint set

final class UnionFind {
    private var root = [Int]()
    private var rank = [Int]()
    
    init(size: Int) {
        rank = [Int](repeating: 1, count: size)
        for i in 0..<size {
            root.append(i)
        }
    }
    
    func find(_ x: Int) -> Int {
        let rootX = root[x]
        if rootX == x {
            return rootX
        }
        root[x] = find(rootX)
        return root[x]
    }
    
    func union(x: Int, y: Int) {
        let rootX = find(x)
        let rootY = find(y)
        if rank[x] > rank[y] {
            root[rootY] = rootX
        } else if rank[y] > rank[x] {
            root[rootX] = rootY
        } else {
            root[rootY] = rootX
            rank[rootX] += 1
        }
    }
    
    func connected(x: Int, y: Int) -> Bool {
        find(x) == find(y)
    }
}