Skip to content

Instantly share code, notes, and snippets.

@JackMorris JackMorris/DisjointSet.swift Secret
Created May 3, 2020

Embed
What would you like to do?
/// Maintains a number of sets, each containing a number of
/// `Element`s.
///
/// Supports union-ing two sets (using an `Element` contained in
/// each set), as well as determining whether two `Element`s are
/// contained within the same set.
///
/// Note that this is a reference type, rather than a struct
/// with value semantics. This is due to the complexity of
/// copying the `Node`s on write, which whilst possible,
/// detracts from the actual behaviour of the type in this
/// example.
public final class DisjointSet<Element: Hashable> {
// MARK: - Types
/// `Error` thrown when an operation expects than an element
/// has already been inserted with `insert(_:)`.
public struct MissingElementError: Error {}
/// A single node in a set.
///
/// Note that this is a reference type, since each `Node` will
/// be stored in the `nodes` dictionary (to lookup `Node`s by
/// an `Element`), and will be referred to by other `Node`s in
/// its set.
private final class Node {
/// The parent `Node`.
///
/// Will be `nil` if this `Node` is the root node of a set.
var parent: Node?
/// An upper-bound on the number of child nodes of this
/// `Node`.
var rank: Int = 0
}
// MARK: - Properties
/// A mapping of `Element`s to `Node`s, containing a `Node`
/// for each `Element` that has been inserted by `insert(_:)`.
private var nodes: [Element: Node] = [:]
// MARK: - Initializers
/// Initializes a `DisjointSet` that initially contains 0
/// elements.
public init() {}
// MARK: - Public
/// Inserts the specified `element`.
///
/// - Returns: `true` if `element` was inserted, or `false`
/// otherwise (e.g. if `element` has already been inserted).
@discardableResult public func insert(
_ element: Element
) -> Bool {
guard nodes[element] == nil else {
// Element already inserted.
return false
}
nodes[element] = Node()
return true
}
/// Unions the two sets containing `lhs` and `rhs`.
///
/// - Parameters:
/// - lhs: An element contained in the first set.
/// - rhs: An element contained in the second set.
/// - Throws: A `MissingElementError` if either `lhs` or `rhs`
/// has not previously been inserted.
public func union(_ lhs: Element, _ rhs: Element) throws {
guard
let lhsNode = nodes[lhs],
let rhsNode = nodes[rhs]
else {
throw MissingElementError()
}
let lhsRoot = findRoot(node: lhsNode)
let rhsRoot = findRoot(node: rhsNode)
guard lhsRoot !== rhsRoot else {
// `lhs` and `rhs` are already in the same set.
return
}
// Union by rank heuristic.
if lhsRoot.rank > rhsRoot.rank {
// Note: `lhsRoot.rank` is intentionally not incremented
// here. `lhsRoot.rank` is still valid, since it is just
// an upper bound (and the new child, `rhsRoot`, has a
// lower rank).
rhsRoot.parent = lhsRoot
} else if lhsRoot.rank < rhsRoot.rank {
// Note: `rhsRoot.rank` is intentionally not incremented
// here (same reasoning as above).
lhsRoot.parent = rhsRoot
} else {
// `lhsRoot.rank == rhsRoot.rank`. Use `rhsRoot` as the
// new root arbitrarily.
lhsRoot.parent = rhsRoot
// `rhsRoot.rank` needs to be updated here to ensure that
// the upper bound is still valid.
rhsRoot.rank += 1
}
}
/// Returns `true` if `lhs` and `rhs` are members of the same
/// set.
///
/// - Parameters:
/// - lhs: The first `Element` to check for co-membership.
/// - rhs: The second `Element` to check for co-membership.
/// - Throws: `true` if `lhs` and `rhs` are contained in the
/// same set (as a result of a `union(_:_:)` operation.
/// - Returns: A `MissingElementError` if either `lhs` or
/// `rhs` has not previously been inserted.
public func sameSet(
_ lhs: Element,
_ rhs: Element
) throws -> Bool {
guard
let lhsNode = nodes[lhs],
let rhsNode = nodes[rhs]
else {
throw MissingElementError()
}
/// Use `===` to check for instance equality, rather than
/// value equality.
return findRoot(node: lhsNode) === findRoot(node: rhsNode)
}
// MARK: - Private
/// Returns the root `Node` pointed to be `node`, or `node`
/// itself if it is a root node.
private func findRoot(node: Node) -> Node {
if let parent = node.parent {
// Path compression heuristic.
node.parent = findRoot(node: parent)
}
// If `node.parent` is `nil`, `node` is a root node.
// Therefore, return it instead.
return node.parent ?? node
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.