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

/* Graphs and MSTs will not be listed if number of nodes exceeds this value. */
#define OUTPUT_THRESHOLD 13

#define NIL        (-1)
#define INFINITY   99999
#define MAXNODES   1000
#define MAXEDGES   (MAXNODES*MAXNODES)
#define MAXWEIGHT  100
#define TRUE       1
#define FALSE      0
#define WHITE      1
#define BLACK      0

/* These data structures are for use in Kruskal's algorithm. */

struct edge {
  int from;
  int to;
  int weight;
};

struct vertex {
  int rank;
  int parent;
};


/*********GRAPH GENERATION*******************************************/


/* minimal DFS -- color only; no tree and no timestamps */
void dfs_visit(int graph[MAXNODES][MAXNODES], int color[MAXNODES], int n, int u) {
  int i;
  color[u] = BLACK;
  for (i=0; i < n; i++) {
    if (graph[u][i] == NIL) continue;
    if (color[i] == WHITE) {
      dfs_visit(graph, color, n, i);
    }
  }
}


/* n m */
void generate_graph(int graph[MAXNODES][MAXNODES], int n, int m)
{
  int count, u, v;
  int i, j;
  int connected;
  int color[MAXNODES];

  connected = FALSE;
  while (! connected) {

    /* clear graph */
    for (i=0; i < n; i++) {
      for (j=0; j < n; j++) {
        graph[i][j] = NIL;
      }
    }

    /* add random edges */
    count = 0;
    srand((int)time(NULL));
    while (count < m) {
      u = rand() % n;
      v = rand() % n;
      if (graph[u][v] == NIL) {
        graph[u][v] = 1 + (rand() % MAXWEIGHT);
        graph[v][u] = graph[u][v];
      }
      ++count;
    }

    /* check by depth-first search whether connected */
    connected = TRUE;
    for (i=0; i < n; i++)
      color[i] = WHITE;
    dfs_visit(graph, color, n, 0);
    /* if there are any unreachable nodes then try again */
    for (i=0; i < n; i++) {
      if (color[i] == WHITE) {
        connected = FALSE;
        break;
      }
    }

  }
}

/* converts adjacency matrix to an array of struct edge. */
/* returns actual number of edges. */
int extract_edge_list(int graph[MAXNODES][MAXNODES], struct edge **edges, int n) {
  int i, j, m;
  struct edge *e;
  m = 0;
  for (i=0; i < n; i++) {
    for (j=0; j < i; j++) {
      if (graph[i][j] == NIL) continue;
      e = (struct edge *)malloc(sizeof(struct edge));
      e->from = i;
      e->to = j;
      e->weight = graph[i][j];
      edges[m] = e;
      ++m;
    }
  }
  return m;
}


/*********GRAPH/MST OUTPUT*******************************************/


/* list the graph edges */
/* returns number of edges */
int print_graph(FILE *out, int graph[MAXNODES][MAXNODES], int n)
{
  int i, j, count;
  count = 0;
  for (i=0; i < n; i++) {
    for (j=0; j < i; j++) {
      if (graph[i][j] == NIL) continue;
      ++count;
      fprintf(out, "n%d -- n%d [label=\"%d\"];\n", i, j, graph[i][j]);
    }
  }
  return count;
}


/* list the MST edges (from predecessor array) */
/* returns total weight of MST */
int print_mst_prim(FILE *out, int graph[MAXNODES][MAXNODES], int pred[MAXNODES], int n) {
  int i, sum;
  sum = 0;
  for (i=0; i < n; i++) {
    if (pred[i] == NIL) continue;
    fprintf(out, "n%d -- n%d [label=\"%d\"];\n", i, pred[i], graph[i][pred[i]]);
    sum += graph[i][pred[i]];
  }
  return sum;
}


/* list the MST edges (from edge list) */
/* returns total weight of MST */
int print_mst_kruskal(FILE *out, struct edge **mst, int n) {
  int i, sum;
  sum = 0;
  for (i=0; i < n - 1; i++) {
    fprintf(out, "n%d -- n%d [label=\"%d\"];\n", mst[i]->from, mst[i]->to, mst[i]->weight);
    sum += mst[i]->weight;
  }
  return sum;
}


/* Create a postscript file that graphically represents
 * the graph.  Uses neato, which is part of the graphviz package.
 * graphviz can be obtained from http://www.research.att.com/sw/tools/graphviz/
 */
/* The mode parameter must be one of "graph", "prim" or "kruskal". */
void graph_diagram(int graph[MAXNODES][MAXNODES], int pred[MAXNODES], struct edge **mst, int n, char *mode) {
  FILE *neatofile;
  char command[80];

  sprintf(command, "tmp_%s.neato", mode);
  neatofile = fopen(command, "w");
  
  fprintf(neatofile, "graph \"graph\" {\n");
  fprintf(neatofile, "nodesep=3;\n");
  fprintf(neatofile, "node [fontsize=10,height=0.3,width=0.3];\n");
  fprintf(neatofile, "edge [fontsize=9, labeldistance]\n");
  if (strcmp(mode, "graph") != 0) {
    fprintf(neatofile, "edge [style=invis,color=gray];\n");
  }
  print_graph(neatofile, graph, n);
  fprintf(neatofile, "edge [style=solid,color=black];\n");
  if (strcmp(mode, "prim") == 0) {
    print_mst_prim(neatofile, graph, pred, n);
  }
  if (strcmp(mode, "kruskal") == 0) {
    print_mst_kruskal(neatofile, mst, n);
  }
  fprintf(neatofile, "}\n");

  fclose(neatofile);
  sprintf(command, "neato -Tps tmp_%s.neato -o tmp_%s.ps", mode, mode);
  system(command);
  printf("(wrote graph diagram to tmp_%s.ps)\n", mode);
}


/*********PRIM'S ALGORITHM*******************************************/


/* linear-time extraction of minimum-key index (must have Q flag) */
int extract_min(int Q[MAXNODES], int key[MAXNODES], int n) {
  int i, x, min;
  min = INFINITY;
  x = NIL;
  for (i=0; i < n; i++) {
    if (! Q[i]) continue;
    if (key[i] < min) {
      min = key[i];
      x = i;
    }
  }
  if (x == NIL) return NIL;
  Q[x] = FALSE;
  return x;
}

/* find an MST in graph -- Prim's algorithm */
void mst_prim(int graph[MAXNODES][MAXNODES], int pred[MAXNODES], int n)
{
  int key[MAXNODES], Q[MAXNODES];
  int i, count;
  int u, v, w;
  /* initialisation */
  for (i=0; i < n; i++) {
    pred[i] = NIL;
    key[i] = INFINITY;
    Q[i] = TRUE;
  }
  key[0] = 0;
  count = 0;
  /* MST generation */
  while (count < n) {
    u = extract_min(Q, key, n);
    for (v=0; v < n; v++) {
      w = graph[u][v];
      if (w == NIL) continue;
      if (Q[v] && w < key[v]) {
        pred[v] = u;
        key[v] = w;
      }
    }
    ++count;
  }
}


/*********KRUSKAL'S ALGORITHM****************************************/


/* core of the quicksort algorithm */
int partition(struct edge **edges, int p, int r) {
  int i, j, w;
  struct edge *e_tmp;
  w = edges[r]->weight;
  i = p-1;
  for (j=p; j < r; j++) {
    if (edges[j]->weight <= w) {
      i = i + 1;
      e_tmp = edges[i];
      edges[i] = edges[j];
      edges[j] = e_tmp;
    }
  }
  e_tmp = edges[i+1];
  edges[i+1] = edges[r];
  edges[r] = e_tmp;
  return i+1;
}

/* quicksort the edges into nondecreasing order of weight */
void sort_edges(struct edge **edges, int p, int r) {
  int q;
  if (p < r) {
    q = partition(edges, p, r);
    sort_edges(edges, p, q-1);
    sort_edges(edges, q+1, r);
  }
}

/* the following 4 functions implement disjoint set forests. */
/* we use the "union by rank" and "path compression" heuristics. */

void make_set(struct vertex **vertices, int x) {
  struct vertex *v;
  v = (struct vertex *)malloc(sizeof(struct vertex));
  v->parent = x;
  v->rank = 0;
  vertices[x] = v;
}

int find_set(struct vertex **vertices, int x) {
  if (vertices[x]->parent != x) {
    vertices[x]->parent = find_set(vertices, vertices[x]->parent);
  }
  return vertices[x]->parent;
}

void set_link(struct vertex **vertices, int x, int y) {
  int xrank, yrank;
  xrank = vertices[x]->rank;
  yrank = vertices[y]->rank;
  if (xrank > yrank) {
    vertices[y]->parent = x;
  } else {
    vertices[x]->parent = y;
    if (xrank == yrank) {
      vertices[y]->rank = yrank + 1;
    }
  }
}

void set_union(struct vertex **vertices, int x, int y) {
  set_link(vertices, find_set(vertices, x), find_set(vertices, y));
}

/* find an MST in graph -- Kruskal's algorithm */
void mst_kruskal(struct edge **edges, struct edge **mst, int n, int m) {
  int i, u, v;
  int tail;
  struct vertex * vertices[MAXNODES];
  /* initialisation */
  for (i=0; i < n; i++) {
    make_set(vertices, i);
  }
  tail = 0;
  sort_edges(edges, 0, m - 1);
  /* MST generation */
  for (i=0; i < m; i++) {
    u = edges[i]->from;
    v = edges[i]->to;
    if (find_set(vertices, u) != find_set(vertices, v)) {
      mst[tail] = edges[i];
      ++tail;
      set_union(vertices, u, v);
    }
  }
}


/*********MAIN PROGRAM***********************************************/


int main(int argc, char *argv[])
{
  int n, m, m_max;
  int weight;
  double prim_time, kruskal_time;
  double reftime;

  /* For Prim's Algorithm: */
  /* Graph represented by an adjacency matrix */
  int graph[MAXNODES][MAXNODES];
  /* MST represented by a predecessor array */
  int pred[MAXNODES];

  /* For Kruskal's Algorithm: */
  /* Graph represented by an array of edges */
  struct edge * edges[MAXEDGES];
  /* MST represented by an array of edges */
  struct edge * mst_edges[MAXNODES];


  if ((argc < 3)
      || (sscanf(argv[1], "%u", &n) != 1)
      || (sscanf(argv[2], "%u", &m_max) != 1)) {
    fprintf(stderr, "Usage: mst n m [x]\nThe graph will have n vertices and max m edges.\n");
    fprintf(stderr, "An optional third argument will trigger graph diagram generation.\n");
    exit(1);
  }
  if (n > MAXNODES) {
    fprintf(stderr, "The number of nodes exceeds the maximum (%d). Recompile.\n", MAXNODES);
    exit(1);
  }


  generate_graph(graph, n, m_max);
  m = extract_edge_list(graph, edges, n);

  /* calculate the MSTs (and time it) */

  reftime = (double)clock();
  mst_prim(graph, pred, n);
  prim_time = (double)clock() - reftime;

  reftime = (double)clock();
  mst_kruskal(edges, mst_edges, n, m);
  kruskal_time = (double)clock() - reftime;

  /* print out graph and MSTs */

  if (n <= OUTPUT_THRESHOLD) {

    printf("Graph edge listing:\n");
    weight = print_graph(stdout, graph, n);
    printf("Number of edges is %d\n\n", weight);

    printf("MST edges via Prim's algorithm:\n");
    weight = print_mst_prim(stdout, graph, pred, n);
    printf("Total weight is %d\n\n", weight);

    printf("MST edges via Kruskal's algorithm:\n");
    weight = print_mst_kruskal(stdout, mst_edges, n);
    printf("Total weight is %d\n\n", weight);

  }

  printf("Computation time was %1.4f seconds for Prim's algorithm.\n", prim_time / (double)CLOCKS_PER_SEC);
  printf("Computation time was %1.4f seconds for Kruskal's algorithm.\n", kruskal_time / (double)CLOCKS_PER_SEC);
  

  if (argc > 3) {
    graph_diagram(graph, pred, mst_edges, n, "graph");
    graph_diagram(graph, pred, mst_edges, n, "prim");
    graph_diagram(graph, pred, mst_edges, n, "kruskal");
  }

  return 0;
}