/// 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