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

typedef struct {            /* A point in the x-y plane */
   double x;                /* x-coordinate */
   double y;                /* y-coordinate */
} location;

typedef struct _node {      /* A node in the adjacency list */
   int nbr;                 /* Id of neighbor */
   struct _node *next;      /* The next link */
} node;

typedef node *graph;        /* A graph is a dynamic array of nodes */

typedef int *partition;     /* Partition array to store 1 or 2 in each cell */

typedef struct {            /* Status of neighbors of a node */
   int nbr1;                /* Number of neighbors in the same part */
   int nbr2;                /* Number of neighbors in the other part */
} nbrstat;

/* Generate the locations of the n houses randomly in the unit square */
location *genRandLocation ( int n )
{
   int i;
   location *loc;

   printf("Houses are located at:\n");
   loc = (location *)malloc(n * sizeof(location));
   for (i=0; i<n; ++i) {
      loc[i].x = (double)rand() / (double)RAND_MAX;
      loc[i].y = (double)rand() / (double)RAND_MAX;
      printf("(%6.4lf,%6.4lf) ", loc[i].x, loc[i].y);
      if (i % 5 == 4) printf("\n");
   }
   if (i % 5 != 0) printf("\n");
   return loc;
}

/* Euclidean distance between two points */
double distance ( location P1, location P2 )
{
   double deltax, deltay;

   deltax = (P2.x - P1.x);
   deltay = (P2.y - P1.y);
   return sqrt(deltax*deltax + deltay*deltay);
}

/* Given the location array for n cities, generate the graph of unsafe pairs */
graph genGraph ( location *loc, int n, double delta )
{
   int i, j;
   graph G;
   node *p;

   G = (node *)malloc(n * sizeof(node));
   for (i=0; i<n; ++i) G[i].next = NULL;
   for (i=0; i<n; ++i) {
      for (j=i+1; j<n; ++j) {
         if (distance(loc[i],loc[j]) < delta) {
            /* Insert j as a neighbor of i */
            p = (node *)malloc(sizeof(node));
            p -> nbr = j;
            p -> next = G[i].next;
            G[i].next = p;

            /* Insert i as a neighbor of j */
            p = (node *)malloc(sizeof(node));
            p -> nbr = i;
            p -> next = G[j].next;
            G[j].next = p;
         }
      }
   }
   return G;
}

/* Find the count of neighbors of u in the same part and also the count
   of neighbors in the other part (according to partition P) */
nbrstat getNbrStat ( graph G, partition P, int u )
{
   node *p;
   nbrstat ns;

   ns.nbr1 = ns.nbr2 = 0;
   p = G[u].next;
   while (p != NULL) { /* p runs over all neighbors in the adjacency list */
      if (P[p -> nbr] == P[u])
         ++ns.nbr1;    /* Neighbor in the same part as u */
     else
         ++ns.nbr2;    /* Neighbor in the other part */
      p = p -> next;
   }
   return ns;
}

/* Calculate the total number of edges connecting pairs of vertices
   in the two parts */
int numberCrossEdges ( graph G, partition P, int n )
{
   int u, cnt;
   nbrstat ns;

   cnt = 0;
   for (u=0; u<n; ++u) {
      ns = getNbrStat(G,P,u); /* Get local and remote neighbors of u */
      cnt += ns.nbr2;         /* The remote neighbors account for cross edges */
   }
   cnt /= 2;                  /* Each cross edge is counted twice in the loop */
   return cnt;
}

/* Initialize to any partition */
partition initPartition ( graph G, int n1, int n2 )
{
   int i;
   partition P;

   P = (partition)malloc((n1 + n2) * sizeof(int));
   for (i=0; i<n1; ++i) P[i] = 1;     /* The first n1 houses belong to Part 1 */
   for (i=n1; i<n1+n2; ++i) P[i] = 2; /* The last n2 houses belong to Part 2 */
   printf("\tNumber of cross edges = %d\n", numberCrossEdges(G,P,n1+n2));
   return P;
}

/* This is another initial partitioning strategy */
partition randPartition ( graph G, int n1, int n2 )
{
   int i, t;
   partition P;

   P = (partition)malloc((n1 + n2) * sizeof(int));
   for (i=0; i<n1+n2; ++i) P[i] = 1;  /* Initially all houses go to Part 1 */
   for (i=0; i<n2; ++i) {             /* Switch n2 houses randomly to Part 2 */
      do { t = rand() % (n1+n2); } while (P[t] == 2);
      P[t] = 2;
   }
   printf("\tNumber of cross edges = %d\n", numberCrossEdges(G,P,n1+n2));
   return P;
}

/* Check whether u and v are neighbors, that is, houses u and v are within
   the safety distance */
int isNbr ( graph G, int u, int v )
{
   node *p;

   p = G[u].next;
   while (p != NULL) {   /* p runs over the adjacency list for u */
      if (p -> nbr == v) return 1;
      p = p -> next;
   }
   return 0;
}

/* Refine partition P:
   Return 1 if any positive refinement is made.
   Return 0 if no further refinement is possible. */
int refinePartition ( partition P, graph G, int n )
{
   nbrstat ns;
   int maxref1, maxref2, thisref;
   int i, u, v;

   maxref1 = maxref2 = -n;

   /* Locate the vertex u in Part 1 such that switching u to Part 2
      reduces the maximum number of cross edges */
   for (i=0; i<n; ++i) {
      if (P[i] == 1) {
         ns = getNbrStat(G,P,i);
         thisref = ns.nbr2 - ns.nbr1;
         if (thisref > maxref1) {
            maxref1 = thisref;
            u = i;
         }
      }
   }

   /* Locate the vertex v in Part 1 such that switching v to Part 1
      reduces the maximum number of cross edges. This v should not be
      a neighbor of u chosen above. */
   for (i=0; i<n; ++i) {
      if ((P[i] == 2) && (!isNbr(G,i,u))) {
         ns = getNbrStat(G,P,i);
         thisref = ns.nbr2 - ns.nbr1;
         if (thisref > maxref2) {
            maxref2 = thisref;
            v = i;
         }
      }
   }

   /* If the two switches give positive benefit, accept it. */
   if (maxref1 + maxref2 > 0) {
      printf("\t   +++ Switching node %d to Part 2: Gain = %d\n", u, maxref1);
      printf("\t   +++ Switching node %d to Part 1: Gain = %d\n", v, maxref2);
      P[u] = 2; P[v] = 1;
      printf("\tNumber of cross edges = %d\n", numberCrossEdges(G,P,n));
      return 1;
   }
   printf("\t   +++ No further refinement possible\n");
   return 0;
}

/* Print which houses are given to Tribe 1 families and which to Tribe 2 families */
void printPartition ( partition P, int n )
{
   int i, cnt;

   printf("Tribe 1 families get the houses:\n");
   cnt = 0;
   for (i=0; i<n; ++i) if (P[i] == 1) {
      printf("%5d", i);
      ++cnt;
      if (cnt % 16 == 0) printf("\n");
   }
   if (cnt % 16) printf("\n");
   printf("Total number of houses = %d\n", cnt);

   printf("Tribe 2 families get the houses:\n");
   cnt = 0;
   for (i=0; i<n; ++i) if (P[i] == 2) {
      printf("%5d", i);
      ++cnt;
      if (cnt % 16 == 0) printf("\n");
   }
   if (cnt % 16) printf("\n");
   printf("Total number of houses = %d\n", cnt);
}

/* Free all dynamically allocated memory */
void windUp ( location *loc, graph G, partition P, int n )
{
   node *p, *q;
   int i;

   for (i=0; i<n; ++i) {
      q = G[i].next;
      while (q != NULL) { /* Free each node in the adjacency list */
         p = q;
         q = q -> next;
         free(p);
      }
   }
   free(loc); free(P); free(G);
}

int main ( int argc, char *argv[] )
{
   int n, n1, n2;
   double delta;
   location *loc;
   graph G;
   partition P;

   srand((unsigned int)time(NULL));
   if (argc == 4) {
      n1 = atoi(argv[1]);
      n2 = atoi(argv[2]);
      delta = atof(argv[3]);
   } else {
      printf("n1 = "); scanf("%d", &n1);
      printf("n2 = "); scanf("%d", &n2);
      printf("delta = "); scanf("%lf", &delta);
   }
   n = n1 + n2;
   loc = genRandLocation(n);
   G = genGraph(loc,n,delta);

   printf("\nTrying first partitioing strategy\n");
   P = initPartition(G,n1,n2);
   while (refinePartition(P,G,n));
   printPartition(P,n);

   printf("\nTrying random partitioing strategy\n");
   P = randPartition(G,n1,n2);
   while (refinePartition(P,G,n));
   printPartition(P,n);

   windUp(loc,G,P,n);
   exit(0);
}