Skip to content

Instantly share code, notes, and snippets.

@V8tr
Created February 16, 2022 16:43
Show Gist options
  • Save V8tr/61086a509fd141c5c3de30a85e85fd90 to your computer and use it in GitHub Desktop.
Save V8tr/61086a509fd141c5c3de30a85e85fd90 to your computer and use it in GitHub Desktop.
Union find with path compression and weighted union
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
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment