/***********************************************************/
/**** C source code for Euclidean minimum spanning tree ****/
/****       Created by Abhijit Das on 29-Oct-2012       ****/
/***********************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>

#include "EMST.h"

/* Create an m x m hash table of n randomly generated points */
hashTable genRandomPoints ( int n )
{
   int i, j, k;
   double x, y;
   hashTable H;
   hashNodePointer p;

   /* Set sizes */
   H.n = n; H.m = (int)sqrt(n); if (n != H.m * H.m) ++H.m;

   /* Allocate space for x- and y-coordinates of points */
   H.x = (double *)malloc(n * sizeof(double));
   H.y = (double *)malloc(n * sizeof(double));

   /* Create an m x m array of list headers */
   H.hdr = (hashNodePointer **)malloc(H.m * sizeof(hashNodePointer *));
   for (i=0; i<H.m; ++i) {
      H.hdr[i] = (hashNodePointer *)malloc(H.m * sizeof(hashNodePointer));
      for (j=0; j<H.m; ++j) H.hdr[i][j] = NULL;
   }

   for (k=0; k<n; ++k) {
      /* Generate the x and y coordinates randomly in [0,1) */
      x = (double)rand() / ((double)RAND_MAX + 1);
      y = (double)rand() / ((double)RAND_MAX + 1);

      /*Locate the cell for (x,y) */
      i = (int)(x * (double)H.m); j = (int)(y * (double)H.m);

      /* Insert (x,y) in the (i,j)-th list */
      p = (hashNode *)malloc(sizeof(hashNode));
      p -> v = k; p -> next = H.hdr[i][j]; H.hdr[i][j] = p;

      /* Store the coordinates of the k-th point */
      H.x[k] = x; H.y[k] = y;
   }

   return H;
}

void freeHashTable ( hashTable H )
{
   int i, j;
   hashNodePointer p, q;

   for (i=0; i<H.m; ++i) {
      for (j=0; j<H.m; ++j) {
         /* Free the (i,j)-th list node by node */
         p = H.hdr[i][j];
         while (p != NULL) {
            q = p -> next;
            free(p);
            p = q;
         }
      }
      free(H.hdr[i]);
   }
   free(H.hdr); free(H.x); free(H.y);
}

/* Given the hash table H and a distance d, create the neighborhood graph.
   This graph would store enough information so that the hash table is not
   needed in the call of Kruskal(). */
graph createNbdGraph ( hashTable H, double d )
{
   graph G;
   int u, v, i, j, imin, imax, jmin, jmax;
   adjNode *p;
   hashNode *q;
   double distance;

   /* Initialize vertex and edge counts */
   G.n = H.n; G.e = 0;

   /* Initialize adjacency-list headers */
   G.E = (adjNode **)malloc(G.n * sizeof(adjNode *));
   for (i=0; i<G.n; ++i) G.E[i] = NULL;

   /* Create the adjacency list for each vertex */
   for (u=0; u<G.n; ++u) {
      /* Determine the search box for u */
      imin = (int)((H.x[u] - d) * (double)H.m); if (imin < 0) imin = 0;
      imax = (int)((H.x[u] + d) * (double)H.m); if (imax >= H.m) imax = H.m-1;
      jmin = (int)((H.y[u] - d) * (double)H.m); if (jmin < 0) jmin = 0;
      jmax = (int)((H.y[u] + d) * (double)H.m); if (jmax >= H.m) jmax = H.m-1;

      /* Locate all points in the search box */
      for (i=imin; i<=imax; ++i) {
         for (j=jmin; j<=jmax; ++j) {
            q = H.hdr[i][j];
            while (q != NULL) {
               /* Locate a neighbor of u */
               v = q -> v;

               /* Consider each pair (u,v) only once */
               if (u < v) {
                  /* Calculate the actual distance between u and v */
                  distance = (H.x[u] - H.x[v]) * (H.x[u] - H.x[v]);
                  distance += (H.y[u] - H.y[v]) * (H.y[u] - H.y[v]);
                  distance = sqrt(distance);

                  /* If the distance is no more than the allowed distance d */
                  if (distance <= d) {
                     /* Insert v in the adjaccency list for u */
                     p = (adjNode *)malloc(sizeof(adjNode));
                     p -> v = v; p -> cost = distance;
                     p -> next = G.E[u]; G.E[u] = p;

                     /* Uncomment the following three lines, if you want to
                        insert u in the adjaccency list for v. This is wasteful
                        in terms of space and is not necessary for Kruskal's
                        algorithm. */
                     /* p = (adjNode *)malloc(sizeof(adjNode));
                        p -> v = u; p -> cost = distance;
                        p -> next = G.E[v]; G.E[v] = p; */

                     /* Increase edge count */
                     ++G.e;
                  }
               }
               q = q -> next;
            }
         }
      }
   }

   return G;
}

/* Add new edges to the graph. The hash table H and the earlier graph G are
   needed. The previous neighborhood distance was d, so we do not reinsert
   pairs which are less than (or equal to) the distance d apart. The new
   neighborhood distance is r. We add new neighbors with respect to the
   distance r only for the t vertices supplied in the array V. */
graph augmentNbdGraph ( hashTable H, graph G, double r, double d, vertex *V, int t )
{
   int u, v, i, j, k, imin, imax, jmin, jmax;
   adjNode *p;
   hashNode *q;
   double distance;

   /* There is no need to reinitialize the earlier graph G */

   /* Augment the adjacency list for each vertex specified in V */
   for (k=0; k<t; ++k) {
      u = V[k];

      /* Determine the search box for u */
      imin = (int)((H.x[u] - r) * (double)H.m); if (imin < 0) imin = 0;
      imax = (int)((H.x[u] + r) * (double)H.m); if (imax >= H.m) imax = H.m-1;
      jmin = (int)((H.y[u] - r) * (double)H.m); if (jmin < 0) jmin = 0;
      jmax = (int)((H.y[u] + r) * (double)H.m); if (jmax >= H.m) jmax = H.m-1;

      /* Locate all points in the search box */
      for (i=imin; i<=imax; ++i) {
         for (j=jmin; j<=jmax; ++j) {
            q = H.hdr[i][j];
            while (q != NULL) {
               /* Locate a neighbor of u */
               v = q -> v;

               /* Consider each pair (u,v) only once */
               if (u < v) {
                  /* Calculate the actual distance between u and v */
                  distance = (H.x[u] - H.x[v]) * (H.x[u] - H.x[v]);
                  distance += (H.y[u] - H.y[v]) * (H.y[u] - H.y[v]);
                  distance = sqrt(distance);

                  /* If the distance is no more than the allowed distance r,
                     but larger than the previous distance d */
                  if ((distance > d) && (distance <= r)) {
                     /* Insert v in the adjaccency list for u */
                     p = (adjNode *)malloc(sizeof(adjNode));
                     p -> v = v; p -> cost = distance;
                     p -> next = G.E[u]; G.E[u] = p;

                     /* Uncomment the following three lines, if you want to
                        insert u in the adjaccency list for v. This is wasteful
                        in terms of space and is not necessary for Kruskal's
                        algorithm. */
                     /* p = (adjNode *)malloc(sizeof(adjNode));
                        p -> v = u; p -> cost = distance;
                        p -> next = G.E[v]; G.E[v] = p; */

                     /* Increase edge count */
                     ++G.e;
                  }
               }
               q = q -> next;
            }
         }
      }
   }

   return G;
}

void freeGraph ( graph G )
{
   int i;
   adjNode *p, *q;

   for (i=0; i<G.n; ++i) {
      p = G.E[i];
      while (p != NULL) {
         q = p -> next;
         free(p);
         p = q;
      }
   }
   free(G.E);
}

/* Implement standard min-heap functions */

/* Store the edges in G in an array */
heap initHeap ( graph G )
{
    heap Q;
    vertex u, v;
    int k;
    adjNode *p;

    /* Q would store all the edges of the graph G */
    Q.size = G.e;
    Q.E = (edge *)malloc(Q.size * sizeof(edge));

    k = 0; /* The writing index in Q.E */

    /* For each vertex of G */
    for (u=0; u<G.n; ++u) {
       /* Look at the neighbor list of u in G */
       p = G.E[u];

       while (p != NULL) {
          /* v is a neighbor of u */
          v = p -> v;

          /* Avoid duplicate insertion of edges in the heap */
          if (u < v) {
             Q.E[k].u = u; Q.E[k].v = v; Q.E[k].cost = p -> cost;
             ++k;
          } else {
             printf("You might have inserted duplicate edges...\n");
          }

          p = p -> next;
       }
    }

   return Q;
}

void heapify ( edge *E, int n, int i )
{
   int l, r, m;
   edge e;

   while (1) {
      l = 2*i + 1; r = 2*i + 2;
      if (l >= n) break;
      if (l == n-1) m = l; else m = (E[l].cost < E[r].cost) ? l : r;
      if (E[i].cost <= E[m].cost) break;
      e = E[i]; E[i] = E[m]; E[m] = e;
      i = m;
   }
}

heap makeHeap ( heap H )
{
   int i;

   for (i=H.size/2-1; i>=0; --i) heapify(H.E,H.size,i);
   return H;
}

heap deleteMin ( heap H )
{
   H.E[0] = H.E[--H.size];
   heapify(H.E,H.size,0);
   return H;
}

void freeHeap ( heap Q )
{
   free(Q.E);
}

/* Implement standard functions for merge-find sets */

MFSet initMFSet ( int n )
{
   int i;
   MFSet S;

   S = (treeNode **)malloc(n * sizeof(treeNode *));
   for (i=0; i<n; ++i) {
      S[i] = (treeNode *)malloc(sizeof(treeNode));
      S[i] -> size = 1;
      S[i] -> parent = S[i];
   }
   return S;
}

treeNode *findSet ( MFSet S, vertex u )
{
   treeNode *p, *q, *r;

   /* Find */
   r = S[u]; while (r != r -> parent) r = r -> parent;

   /* Do path compression */
   p = S[u];
   while (p -> parent != r) {
      q = p -> parent; p -> parent = r; p = q;
   }

   return r;
}

void mergeSets ( treeNode *r1, treeNode *r2 )
{
   treeNode *r;

   if (r1 -> size <= r2 -> size) r = r1 -> parent = r2;
   else r = r2 -> parent = r1;
   r -> size = r1 -> size + r2 -> size;
}

void freeSet ( MFSet S, int n )
{
   int i;

   for (i=0; i<n; ++i) free(S[i]);
   free(S);
}

/* Find the minimum spanning forest of a graph using Kruskal's algorithm */
SF Kruskal ( graph G, MFSet *S )
{
   SF T;
   heap Q;
   treeNode *r1, *r2;

   /* Initialize the minimum spanning forest */
   T.n = G.n; T.e = 0; T.cost = 0;
   T.E = (edge *)malloc((T.n - 1) * sizeof(edge));

   /* Initialize the heap and the merge-find set */
   Q = makeHeap(initHeap(G));
   *S = initMFSet(T.n);

   /* So long as there are multiple trees and there are more edges waiting
      in the priority queue */
   while ((T.e < T.n-1) && (Q.size > 0)) {
      /* Find the headers r1, r2 of the current minimum-cost edge (u,v) */
      r1 = findSet(*S,Q.E[0].u); r2 = findSet(*S,Q.E[0].v);

      /* If u and v belong to different sets */
      if (r1 != r2) {
         /* Add the edge (u,v) to the forest */
         T.E[T.e].u = Q.E[0].u; T.E[T.e].v = Q.E[0].v;

         /* Add the cost of the edge just inserted in the forest */
         T.cost += (T.E[T.e].cost = Q.E[0].cost);

         /* Increment the count of edges, and merge the two sets of u and v */
         ++T.e; mergeSets(r1,r2);
      }

      /* In any case, we have to delete the minimum from the heap */
      Q = deleteMin(Q);
   }

   /* We no longer need the heap */
   freeHeap(Q);

   return T;
}

void freeSF ( SF T )
{
   free(T.E);
}

/* Return the count and list of vertices outside the largest component */
vertex *leftOverVertices ( SF T, MFSet S, int n )
{
   int i, j, t, *cnt, *V;
   treeNode *r, **R;

   /* The total number of components is computed in t */
   t = T.n - T.e;

   /* Allocate memory for t component headers */
   R = (treeNode **)malloc(t * sizeof(treeNode *));

   /* Also maintain counts of nodes in the t components */
   cnt = (int *)malloc(t * sizeof(int));

   /* Initialize the component headers and counts */
   for (i=0; i<t; ++i) { R[i] = NULL; cnt[i] = 0; }

   /* Locate the component of each vertex */
   for (i=0; i<n; ++i) {
      r = findSet(S,i);

      /* Update information about the component header r */
      for (j=0; j<t; ++j) {
         if (R[j] == r) {              /* r was encountered earlier */
            ++cnt[j]; break;
         } else if (R[j] == NULL) {    /* r is encountered for the first time */
            R[j] = r; cnt[j] = 1; break;
         }
      }
   }

   /* Compute in j the component with largest vertex count. Remember the
      the corresponding component header in r. */
   j = -1;
   for (i=0; i<t; ++i) if (cnt[i] > j) {
      j = cnt[i];
      r = R[i];
   }

   /* We no longer need the arrays R[] and cnt[] */
   free(cnt); free(R);

   /* Make a list of vertices outside the largest component. The number of
      these vertices is n-j. We store this count in V[0], and the vertex
      labels in V[1],V[2],...,V[n-j]. */
   V = (int *)malloc((n-j+1)*sizeof(int));
   V[0] = n-j; j = 0;
   for (i=0; i<n; ++i) if (findSet(S,i) != r) V[++j] = i;

   return V;
}

/* Rajasekaran's EMST algorithm */
SF EMST ( hashTable H )
{
   graph G;
   SF T;
   MFSet S;
   double d, r;
   vertex *V;

   d = C1/sqrt(H.n); r = C2 * log(H.n)/sqrt(H.n);
   printf("\n+++ First phase of EMST...\n");

   /* Create the neighborhood graph with respect to distance d */
   G = createNbdGraph(H,d);
   printf("    No of edges in G is %d\n", G.e);

   /* Find the minimum spanning forest of the initial neighborhood graph */
   T = Kruskal(G,&S);
   printf("    No of edges in T is %d\n", T.e);
   printf("    Number of trees = %d\n", T.n - T.e);

   /* Repeat while the minimum spanning forest contains multiple trees. */
   /* Usually, one iteration of this loop suffices. */
   while (T.e < T.n-1) {
      printf("\n+++ Second phase of EMST...\n");

      /* Find the vertices outside the largest connected component in G */
      V = leftOverVertices(T,S,T.n);
      printf("    Largest tree is of size = %d\n", T.n-V[0]);

      /* Free the MFSet and SF data structures */
      freeSet(S,T.n); freeSF(T);

      /* Do not free the neighborhood graph here. We will keep on adding
         edges to it (instead of rebuilding the graph again from the scratch).
         Indeed, we do not touch the adjacency lists for vertices in the
         largest connected component. */

      /* Add new edges to the remaining vertices with respect to the new distance r */
      G = augmentNbdGraph(H,G,r,d,V+1,V[0]);

      /* V is no longer needed, so free it */
      free(V);
      printf("    No of edges in G is %d\n", G.e);

      /* Find the minimum spanning forest of the augmented neighborhood graph */
      T = Kruskal(G,&S);
      printf("    No of edges in T is %d\n", T.e);
      printf("    Number of trees = %d\n", T.n - T.e);

      /* Prepare for the next iteration */
      d = r; r *= 2;
   }

   /* free dynamically allocated memory not to be used in future */
   freeSet(S,T.n); freeGraph(G);

   return T;
}

/* Print the finally computed MST to a file. Needed for debugging only. */
void printMST ( hashTable H, SF T )
{
   int i;
   FILE *fp;

   fp = (FILE *)fopen("EMST.dat","w");
   fprintf(fp,"n = %d\ne = %d\n", T.n, T.e);
   fprintf(fp,"Location of points\n");
   for (i=0; i<H.n; ++i)
      fprintf(fp, "\t%d (%10.8lf,%10.8lf)\n", i, H.x[i], H.y[i]);
   fprintf(fp,"Edges in MST\n");
   for (i=0; i<T.e; ++i)
      fprintf(fp,"\t(%d,%d) %10.8lf\n", T.E[i].u, T.E[i].v, T.E[i].cost);
   fprintf(fp,"Total cost = %10.8lf\n", T.cost);
   fclose(fp);
}

int main ( int argc, char *argv[] )
{
   int n;
   hashTable H;
   SF T;

   srand((unsigned int)time(NULL));
   if (argc >= 2) n = atoi(argv[1]); else {
      printf("n = "); scanf("%d", &n);
   }

   /* Populate the m x m hash table with n points */
   H = genRandomPoints(n);

   /* Call the EMST algorithm on H */
   T = EMST(H);
   /* printMST(H,T); */

   /* Wind up */
   freeSF(T); freeHashTable(H);

   exit(0);
}

/**** End of EMST.c ****/