Disjoint Set

import DisjointSetItem from './DisjointSetItem';

export default class DisjointSet {
  /**
   * @param {function(value: *)} [keyCallback]
   */
  constructor(keyCallback) {
    this.keyCallback = keyCallback;
    this.items = {};
  }

  /**
   * @param {*} itemValue
   * @return {DisjointSet}
   */
  makeSet(itemValue) {
    const disjointSetItem = new DisjointSetItem(itemValue, this.keyCallback);

    if (!this.items[disjointSetItem.getKey()]) {
      // Add new item only in case if it not presented yet.
      this.items[disjointSetItem.getKey()] = disjointSetItem;
    }

    return this;
  }

  /**
   * Find set representation node.
   *
   * @param {*} itemValue
   * @return {(string|null)}
   */
  find(itemValue) {
    const templateDisjointItem = new DisjointSetItem(itemValue, this.keyCallback);

    // Try to find item itself;
    const requiredDisjointItem = this.items[templateDisjointItem.getKey()];

    if (!requiredDisjointItem) {
      return null;
    }

    return requiredDisjointItem.getRoot().getKey();
  }

  /**
   * Union by rank.
   *
   * @param {*} valueA
   * @param {*} valueB
   * @return {DisjointSet}
   */
  union(valueA, valueB) {
    const rootKeyA = this.find(valueA);
    const rootKeyB = this.find(valueB);

    if (rootKeyA === null || rootKeyB === null) {
      throw new Error('One or two values are not in sets');
    }

    if (rootKeyA === rootKeyB) {
      // In case if both elements are already in the same set then just return its key.
      return this;
    }

    const rootA = this.items[rootKeyA];
    const rootB = this.items[rootKeyB];

    if (rootA.getRank() < rootB.getRank()) {
      // If rootB's tree is bigger then make rootB to be a new root.
      rootB.addChild(rootA);

      return this;
    }

    // If rootA's tree is bigger then make rootA to be a new root.
    rootA.addChild(rootB);

    return this;
  }

  /**
   * @param {*} valueA
   * @param {*} valueB
   * @return {boolean}
   */
  inSameSet(valueA, valueB) {
    const rootKeyA = this.find(valueA);
    const rootKeyB = this.find(valueB);

    if (rootKeyA === null || rootKeyB === null) {
      throw new Error('One or two values are not in sets');
    }

    return rootKeyA === rootKeyB;
  }
}

/**
 * The minimalistic (ad hoc) version of a DisjointSet (or a UnionFind) data structure
 * that doesn't have external dependencies and that is easy to copy-paste and
 * use during the coding interview if allowed by the interviewer (since many
 * data structures in JS are missing).
 *
 * Time Complexity:
 *
 * - Constructor: O(N)
 * - Find: O(α(N))
 * - Union: O(α(N))
 * - Connected: O(α(N))
 *
 * Where N is the number of vertices in the graph.
 * α refers to the Inverse Ackermann function.
 * In practice, we assume it's a constant.
 * In other words, O(α(N)) is regarded as O(1) on average.
 */
class DisjointSetAdhoc {
  /**
   * Initializes the set of specified size.
   * @param {number} size
   */
  constructor(size) {
    // The index of a cell is an id of the node in a set.
    // The value of a cell is an id (index) of the root node.
    // By default, the node is a parent of itself.
    this.roots = new Array(size).fill(0).map((_, i) => i);

    // Using the heights array to record the height of each node.
    // By default each node has a height of 1 because it has no children.
    this.heights = new Array(size).fill(1);
  }

  /**
   * Finds the root of node `a`
   * @param {number} a
   * @returns {number}
   */
  find(a) {
    if (a === this.roots[a]) return a;
    this.roots[a] = this.find(this.roots[a]);
    return this.roots[a];
  }

  /**
   * Joins the `a` and `b` nodes into same set.
   * @param {number} a
   * @param {number} b
   * @returns {number}
   */
  union(a, b) {
    const aRoot = this.find(a);
    const bRoot = this.find(b);

    if (aRoot === bRoot) return;

    if (this.heights[aRoot] > this.heights[bRoot]) {
      this.roots[bRoot] = aRoot;
    } else if (this.heights[aRoot] < this.heights[bRoot]) {
      this.roots[aRoot] = bRoot;
    } else {
      this.roots[bRoot] = aRoot;
      this.heights[aRoot] += 1;
    }
  }

  /**
   * Checks if `a` and `b` belong to the same set.
   * @param {number} a
   * @param {number} b
   */
  connected(a, b) {
    return this.find(a) === this.find(b);
  }
}

export default DisjointSetAdhoc;

export default class DisjointSetItem {
  /**
   * @param {*} value
   * @param {function(value: *)} [keyCallback]
   */
  constructor(value, keyCallback) {
    this.value = value;
    this.keyCallback = keyCallback;
    /** @var {DisjointSetItem} this.parent */
    this.parent = null;
    this.children = {};
  }

  /**
   * @return {*}
   */
  getKey() {
    // Allow user to define custom key generator.
    if (this.keyCallback) {
      return this.keyCallback(this.value);
    }

    // Otherwise use value as a key by default.
    return this.value;
  }

  /**
   * @return {DisjointSetItem}
   */
  getRoot() {
    return this.isRoot() ? this : this.parent.getRoot();
  }

  /**
   * @return {boolean}
   */
  isRoot() {
    return this.parent === null;
  }

  /**
   * Rank basically means the number of all ancestors.
   *
   * @return {number}
   */
  getRank() {
    if (this.getChildren().length === 0) {
      return 0;
    }

    let rank = 0;

    /** @var {DisjointSetItem} child */
    this.getChildren().forEach((child) => {
      // Count child itself.
      rank += 1;

      // Also add all children of current child.
      rank += child.getRank();
    });

    return rank;
  }

  /**
   * @return {DisjointSetItem[]}
   */
  getChildren() {
    return Object.values(this.children);
  }

  /**
   * @param {DisjointSetItem} parentItem
   * @param {boolean} forceSettingParentChild
   * @return {DisjointSetItem}
   */
  setParent(parentItem, forceSettingParentChild = true) {
    this.parent = parentItem;
    if (forceSettingParentChild) {
      parentItem.addChild(this);
    }

    return this;
  }

  /**
   * @param {DisjointSetItem} childItem
   * @return {DisjointSetItem}
   */
  addChild(childItem) {
    this.children[childItem.getKey()] = childItem;
    childItem.setParent(this, false);

    return this;
  }
}