-
-
Save JackMorris/efb55c3ecf9ac7f59cff3a1ef91599c2 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
/// 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