DisjointSet.swift

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

}