/************************************************************************
 * File: UnionFind.java
 * Author: Keith Schwarz (htiek@cs.stanford.edu)
 *
 * An implementation of a union-find (disjoint set) data structure using
 * a disjoint-set forest.  This implementation implements the path
 * compression and union by rank optimizations, which results in very
 * efficient amortized cost for any sequence of operations.  In
 * particular, any m operations will complete in O(m a(m)), where a(m)
 * is the Ackermann inverse function.  This function grows incredibly
 * slowly, and is effectively a constant for any value of m less than
 * the number of atoms in the universe.
 *
 * Disjoint set structures can be used to implement relational
 * unification, Kruskal's MST algorithm, or Hindley-Milner type
 * inference.  They are also good at finding connected components of
 * an undirected graph.
 */


import java.util.*; // For Map, HashMap

/**
 * A class representing the union-find abstraction.
 *
 * @param T The type of elements to store in the structure.
 * @author Keith Schwarz (htiek@cs.stanford.edu)
 */

public final class UnionFind<T> {
    /**
     * A utility struct holding an an object's parent and rank.
     */

    private static final class Link<T> {
        public T parent;
        public int rank = 0;

        /**
         * Creates a new Link object with the specified parent.
         *
         * @param parent The initial value of the parent field.
         */

        Link(T parent) {
            this.parent = parent;
        }
    }

    /**
     * A map from objects in the UnionFind structure to their associated
     * rank and parent.
     */

    private final Map<T, Link<T>> elems = new HashMap<T, Link<T>>();

    /**
     * Creates a new UnionFind structure that is initially empty.
     */

    public UnionFind() {
        // Handled implicitly
    }

    /**
     * Creates a new UnionFind structure that initially contains all of
     * the elements from the specified container.  Duplicate elements
     * will appear with multiplicity one.
     *
     * @param elems The elements to store in the UnionFind structure.
     */

    public UnionFind(Collection<? extends T> elems) {
        /* Iterate across the collection, adding each to this structure. */
        for (T elem: elems)
            add(elem);
    }

    /**
     * Inserts a new element into the UnionFind structure that initially
     * is in its own partition.  If the element already exists, this
     * function returns false.  Otherwise, it returns true.
     *
     * @param elem The element to insert.
     * @return Whether the element was successfully inserted.
     * @throws NullPointerException If elem is null.
     */

    public boolean add(T elem) {
        /* Check for null. */
        if (elem == null)
            throw new NullPointerException("UnionFind does not support null.");

        /* Check whether this entry exists; fail if it does. */
        if (elems.containsKey(elem))
            return false;

        /* Otherwise add the element as its own parent. */
        elems.put(elem, new Link<T>(elem));
        return true;
    }

    /**
     * Given an element, returns the representative element of the set
     * containing that element.  If the element is not contained in the
     * structure, this method throws a NoSuchElementException.
     *
     * @param elem The element to look up.
     * @return The representative of the set containing the element.
     * @throws NoSuchElementException If the element does not exist.
     */

    public T find(T elem) {
        /* Check whether the element exists; fail if it doesn't. */
        if (!elems.containsKey(elem))
            throw new NoSuchElementException(elem + " is not an element.");

        /* Recursively search the structure and return the result. */
        return recFind(elem);
    }

    /**
     * Given an element which is known to be in the structure, searches
     * for its representative element and returns it, applying path
     * compression at each step.
     *
     * @param elem The element to look up.
     * @return The representative of the set containing it.
     */

    private T recFind(T elem) {
        /* Get the info on this object. */
        Link<T> info = elems.get(elem);

        /* If the element is its own parent, it's the representative of its
         * class and we should say so.
         */

        if (info.parent.equals(elem))
            return elem;

        /* Otherwise, look up the parent of this element, then compress the
         * path.
         */

        info.parent = recFind(info.parent);
        return info.parent;
    }
    
    /**
     * Given two elements, unions together the sets containing those
     * elements.  If either element is not contained in the set,
     * throws a NoSuchElementException.
     *
     * @param one The first element to link.
     * @param two The second element to link.
     * @throws NoSuchElementException If either element does not exist.
     */

    public void union(T one, T two) {       
        /* Get the link info for the parents.  This also handles the exception
         * guarantee.
         */

        Link<T> oneLink = elems.get(find(one));
        Link<T> twoLink = elems.get(find(two));

        /* If these are the same object, we're done. */
        if (oneLink == twoLink) return;

        /* Otherwise, link the two.  We'll do a union-by-rank, where the parent
         * with the lower rank will merge with the parent with higher rank.
         */

        if (oneLink.rank > twoLink.rank) {
            /* Since each parent must link to itself, the value of oneLink.parent
             * is the representative of one.
             */

            twoLink.parent = oneLink.parent;
        } else if (oneLink.rank < twoLink.rank) {
            /* Same logic as above. */
            oneLink.parent = twoLink.parent;
        } else {
            /* Arbitrarily wire one to be the parent of two. */
            twoLink.parent = oneLink.parent;
            
            /* Bump up the representative of one to the next rank. */
            oneLink.rank++;
        }
    }
}