/* Union-Find
 * 15-122 Principles of Imperative Computation, Fall 2010
 * Frank Pfenning
 */

typedef struct ufs* ufs;
struct ufs {
  int size;
  int[] A;			/* \length(A) == size */
};

bool elem(ufs eqs, int i) {
  return 0 <= i && i < eqs->size;
}

bool root(ufs eqs, int i)
{
  return elem(eqs, i) && eqs->A[i] < 0;
}

bool depth_bounded(ufs eqs, int i)
{
  int[] A = eqs->A;
  int size = eqs->size;
  int k = i;
  int d = 1;
  while (elem(eqs, k) && A[k] >= 0 && d <= size) {
    d++;
    k = A[k];
  }
  /* !elem(eqs, k) means index k out of range */
  /* A[k] > -d means path exceeds depth bound */
  /* d > size we must be in loop */
  return elem(eqs, k) && A[k] <= -d && d <= size;
}

bool valid_ufs (ufs eqs)
//@requires eqs->size <= \length(eqs->A);
{
  int size = eqs->size;
  int[] A = eqs->A;
  int i;
  if (!(0 <= eqs->size)) return false;
  for (i = 0; i < size; i++)
    if (!depth_bounded(eqs,i))
      return false;
  return true;
}

ufs singletons (int n)
//@requires 0 <= n;
//@ensures valid_ufs(\result);
{
  int i;
  int[] A = alloc_array(int, n);
  ufs eqs = alloc(struct ufs);
  for (i = 0; i < n; i++) {
    A[i] = -1;
  }
  eqs->size = n;
  eqs->A = A;
  return eqs;
}

int find (ufs eqs, int i)
//@requires valid_ufs(eqs) && elem(eqs,i);
//@ensures valid_ufs(eqs) && root(eqs, \result);
{
  int[] A = eqs->A;
  int k = i;
  while (A[k] >= 0) 
    //@loop_invariant elem(eqs, k);
    k = A[k];
  //@assert root(eqs, k);
  if (i != k) A[i] = k;		/* path compression */
  return k;
}

int union (ufs eqs, int i, int k)
//@requires valid_ufs(eqs);
//@requires elem(eqs, i) && elem(eqs, k);
//@ensures valid_ufs(eqs);
//@ensures find(eqs, i) == find(eqs, k);
{
  int iclass = find(eqs, i);
  int kclass = find(eqs, k);
  int[] A = eqs->A;
  if (iclass == kclass) return iclass;
  if (A[iclass] < A[kclass]) {  /* i has greater depth */
    A[kclass] = iclass;		/* depth remains the same */
    return iclass;
  } else if (A[iclass] == A[kclass]) {
    A[kclass] = iclass;		/* direction is arbitrary */
    A[iclass]--;		/* depth increases by one */
    return iclass;
  } else {			/* k has greater depth */
    A[iclass] = kclass;		/* depth remains the same */
    return kclass;
  }
}