/*******************************************************************\

Module:

Author: Daniel Kroening, kroening@cs.cmu.edu

\*******************************************************************/

#include "union_find.h"

/*******************************************************************\

Function: unsigned_union_find::make_union

  Inputs:

 Outputs:

 Purpose:

\*******************************************************************/

void unsigned_union_find::make_union(unsigned j, unsigned k)
{
  check_index(j);
  check_index(k);

  // make sure j, k are roots
  j=find(j);
  k=find(k);

  if(j==k) return; // already in same set

  // weight it

  if(nodes[j].count < nodes[k].count)  // j is the smaller set
  {
    nodes[k].count+=nodes[j].count;  // so k becomes parent to j
    nodes[j].root=false;
    nodes[j].parent=k;
  }
  else // j is NOT the smaller
  {
    nodes[j].count+=nodes[k].count;  // so j becomes parent to k
    nodes[k].root=false;
    nodes[k].parent=j;
  }
}

/*******************************************************************\

Function: unsigned_union_find::isolate

  Inputs:

 Outputs:

 Purpose:

\*******************************************************************/

void unsigned_union_find::isolate(unsigned a)
{
  check_index(a);
  
  // already isolated?
  if(is_root(a) && nodes[a].count==1)
    return;

  unsigned_union_find new_sets;
  
  new_sets.resize(size());
  
  // is a itself a root?
  
  if(is_root(a))
  {
    unsigned c=nodes[a].count;
    assert(c>=2);

    // find a new root
    unsigned new_root=a;
    
    for(unsigned i=0; i<size(); i++)
      if(find(i)==a && i!=a)
      {
        new_root=i;
        break;
      }
    
    assert(new_root!=a);

    for(unsigned i=0; i<size(); i++)
      if(i!=a && !is_root(i))
      {
        unsigned r=find(i);
        if(r==a) r=new_root;
        new_sets.make_union(i, r);
      }
  }
  else
  {
    for(unsigned i=0; i<size(); i++)
      if(i!=a && !is_root(i))
        new_sets.make_union(i, find(i));
  }
  
  swap(new_sets);
}

/*******************************************************************\

Function: unsigned_union_find::intersection

  Inputs:

 Outputs:

 Purpose:

\*******************************************************************/

void unsigned_union_find::intersection(
  const unsigned_union_find &other)
{
  unsigned_union_find new_sets;
  
  new_sets.resize(std::max(size(), other.size()));

  // should be n log n
  
  for(unsigned i=0; i<size(); i++)
    if(!is_root(i))
    {
      unsigned j=find(i);
      
      if(other.same_set(i, j))
        new_sets.make_union(i, j);
    }
  
  swap(new_sets);
}

/*******************************************************************\

Function: unsigned_union_find::find

  Inputs:

 Outputs:

 Purpose:

\*******************************************************************/

unsigned unsigned_union_find::find(unsigned a) const
{
  if(a>=size()) return a;
  while(!nodes[a].root) a=nodes[a].parent;
  return a;
}
