partition.c

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

#define return_null_if_null(p) \
  do {                         \
    if (NULL == (p)) {         \
      return NULL;             \
    }                          \
  } while (0)

#define partition_find_set partition_find_set_r

/**
 * The weight of each node needs about 2^64 values to avoid integer overflow for
 * our problem domain.
 */
typedef uint64_t weight_t;

typedef struct {
  // Pointer to array of node IDs.
  int *id;
  // Pointer to array of tree weights.
  weight_t *weight;
} partition_t;

// NewPartition creates a new partition.
partition_t *partition_new(size_t n) {
  partition_t *p = malloc(sizeof(partition_t));
  if (NULL != p) {
    memset(p, 0, sizeof(partition_t));

    p->id = malloc((n + 1) * sizeof(int));
    if (NULL == p->id) {
      free(p);
      return NULL;
    }

    p->weight = malloc((n + 1) * sizeof(weight_t));
    if (NULL == p->weight) {
      free(p->id);
      free(p);
      return NULL;
    }

    for (size_t i = 0; i < n; i++) {
      p->id[i] = i;
      p->weight[i] = 1;
    }
  }
  return p;
}

int partition_find_set_1(partition_t *p, int x) {
  // Single-pass point-at-grandparent path compression.
  while (x != p->id[x]) {
    p->id[x] = p->id[p->id[x]];
    x = p->id[x];
  }
  return x;
}

int partition_find_set_2(partition_t *p, int x) {
  // Two-pass point-at-root path compression.
  int i = x;
  while (x != p->id[x]) {
    x = p->id[x];
  }
  // x is now root.
  while (i != p->id[i]) {
    i = p->id[i];
    p->id[i] = x;
  }
  return x;
}

int partition_find_set_r(partition_t *p, int x) {
  // Two-pass recursive point-all-nodes-at-root path compression.
  if (x != p->id[x]) {
    p->id[x] = partition_find_set_r(p, p->id[x]);
  }
  return p->id[x];
}

weight_t partition_weight(partition_t *p, int x) {
  return p->weight[partition_find_set(p, x)];
}

void partition_union(partition_t *p, int x, int y) {
  int a = partition_find_set(p, x);
  int b = partition_find_set(p, y);
  if (p->weight[a] < p->weight[b]) {
    p->id[a] = b;
    p->weight[b] += p->weight[a];
  } else {
    p->id[b] = a;
    p->weight[a] += p->weight[b];
  }
}

bool partition_connected(partition_t *p, int x, int y) {
  return partition_find_set(p, x) == partition_find_set(p, y);
}

void partition_free(partition_t *p) {
  free(p->id);
  free(p->weight);
  free(p);
}

void test() {
  partition_t *p = partition_new(10);
  if (NULL != p) {
    for (int i = 0; i < 10; i++) {
      printf("%d %llu\n", p->id[i], p->weight[i]);
    }
  }
  partition_union(p, 4, 3);
  partition_union(p, 3, 8);
  partition_union(p, 6, 5);
  partition_union(p, 9, 4);
  partition_union(p, 2, 1);
  printf("(0, 7) == false: %d\n", partition_connected(p, 0, 7));
  printf("(8, 9) == true: %d\n", partition_connected(p, 8, 9));
  fflush(stdout);
  partition_union(p, 5, 0);
  partition_union(p, 7, 2);
  partition_union(p, 6, 1);
  partition_union(p, 1, 0);
  printf("(0, 7) == true: %d\n", partition_connected(p, 0, 7));
  fflush(stdout);
  partition_free(p);
}

/* This subroutine will read one line of input from a file.
   A line is defined to be an arbitrary string of characters that is
   terminated by a newline character or an end-of-file.  If the line
   from the file exceeds the length provided by maxlen then the rest
   of the line is read and thrown away.  Any error messages are sent to
   standard error.

   The return value is the length of the line that was read, excluding
   the newline character.  If the read begins at the end-of-file, this
   function returns EOF.


   Warning: this routine assumes that the string termination character
   '\0' can be stored at s[maxlen+1] if needed.  Therefore if you declare
   the string to be

       char new_line[80];

   the function call should look like

       n_read = getline(new_line,79);

   Actually, any value less than 80 will do, just don't use the same
   number as appears in the declaration.  You need to leave one character
   for the null character.
*/
int fgetline(FILE *fp, char *s, int maxlen) {
  int i = 0; /* counts the number of characters read (up to maxlen) */
  char c;
  int ok = 1; /* normally TRUE, if false then the input string was too long */

  while ((c = getc(fp)) != EOF) {
    if (c == '\n') {
      *s = '\0';
      return i;
    } else if (ok) {
      if (i >= maxlen) {
        fprintf(stderr,
                "FGETLINE error: input line longer than %d characters, "
                "truncating line.\n",
                maxlen);
        ok = 0;
      } else {
        *s++ = c;
        i++;
      }
    }
  }

  if (0 == i) {
    return EOF;
  } else {
    return i;
  }
}

#define P_LINE_BUF_SIZ (80)

int main(int argc, char *argv[]) {
  int n = 0;
  int q = 0;
  char c = ' ';
  char line[P_LINE_BUF_SIZ + 1];  // Allocate additional space for the NULL
                                  // terminating character.
  int i = 0;
  int j = 0;
  partition_t *p = NULL;

  test();

  // The problem uses 1-based indexing.
  scanf("%d %d\n", &n, &q);
  p = partition_new(n + 1);
  if (NULL != p) {
    while (q-- > 0) {
      c = getchar();
      getchar();  // discard whitespace.
      switch (c) {
        case 'M':
          fgetline(stdin, line, P_LINE_BUF_SIZ);
          sscanf(line, "%d %d\n", &i, &j);
          if (!partition_connected(p, i, j)) {
            partition_union(p, i, j);
          }
          break;
        case 'Q':
          fgetline(stdin, line, P_LINE_BUF_SIZ);
          sscanf(line, "%d\n", &i);
          printf("%llu\n", partition_weight(p, i));
          break;
        default:
          break;
      }
    }
    partition_free(p);
  }
  return 0;
}