Union find with path compression and weighted union

UnionFind.swift

class UnionFind {
    private var id: [Int]
    private var sizes: [Int]
    private(set) var components: Int
    
    init(_ N: Int) {
        id = (0..<N).map { $0 }
        sizes = Array(repeating: 1, count: N)
        components = N
    }
    
    func union(_ p: Int, _ q: Int) {
        let idP = find(p)
        let idQ = find(q)
        
        if idP == idQ {
            return
        }
        
        // Weighted union
        if sizes[idP] > sizes[idQ] {
            id[idQ] = idP
            sizes[idP] += sizes[idQ]
        } else {
            id[idP] = idQ
            sizes[idQ] += sizes[idP]
        }
        
        components -= 1
    }
    
    func find(_ p: Int) -> Int {
        var p = p
      
        // Chase parent pointers until the root is reached
        while p != id[p] {
            // Path compression
            id[p] = id[id[p]]          
            p = id[p]
        }
        return p
    }
}