/*
 * Copyright (c) 2010, The Broad Institute
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 */
package org.broadinstitute.sting.utils;

import java.util.*;

public class DisjointSet {
    private ItemNode[] nodes;

    public DisjointSet(int numItems) {
        this.nodes = new ItemNode[numItems];
        for (int i = 0; i < numItems; i++)
            this.nodes[i] = new ItemNode(i);
    }

    public int findSet(int itemIndex) {
        // Find itemIndex's root Node:
        ItemNode curNode = nodes[itemIndex];
        while (curNode.parent != null)
          curNode = curNode.parent;
        ItemNode root = curNode;

        // Perform path compression:
        curNode = nodes[itemIndex];
        while (curNode != root) {
          ItemNode next = curNode.parent;
          curNode.parent = root;
          curNode = next;
        }

        return root.itemIndex;
    }

    public boolean inSameSet(int x, int y) {
        return (x == y || nodes[x].parent == nodes[y].parent || findSet(x) == findSet(y));
    }

    public Set<Integer> inSameSetAs(int x, Collection<Integer> testSet) {
        Set<Integer> sameSetInds = new TreeSet<Integer>();

        int xSet = findSet(x);
        for (int t : testSet) {
            if (findSet(t) == xSet)
                sameSetInds.add(t);
        }
        return sameSetInds;
    }

    public void setUnion(int x, int y) {
        link(findSet(x), findSet(y));
    }

    private void link(int x, int y) {
        if (x == y)
            return;

        // Union by rank:
        if (nodes[x].rank > nodes[y].rank) {
          nodes[y].parent = nodes[x];
        }
        else { // nodes[x].rank <= nodes[y].rank
          nodes[x].parent = nodes[y];
          if (nodes[x].rank == nodes[y].rank)
            nodes[y].rank++;
        }
    }

    private class ItemNode {
        private int itemIndex;
        private ItemNode parent;
        private int rank;

        public ItemNode(int itemIndex) {
            this.itemIndex = itemIndex;
            this.parent = null;
            this.rank = 0;
        }
    }
}
