avl_tree.cpp

#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>

struct node {
  int val;
  struct node *lhs, *rhs;
  struct node *par;
  int height;
};

struct node *node_new(int val, struct node *par) {
  struct node *res = (struct node *)malloc(sizeof(struct node));
  res->val = val;
  res->lhs = NULL;
  res->rhs = NULL;
  res->par = par;
  res->height = 1;

  return res;
}
void node_free(struct node *self) {
  while (self != NULL) {
    node_free(self->lhs);
    struct node *next = self->rhs;
    free(self);
    self = next;
  }
}

bool node_contains(struct node *self, int val) {
  while (self != NULL) {
    if (self->val == val) return true;
    else if (self->val < val) self = self->lhs;
    else self = self->rhs;
  }

  return false;
}
void node_update_height(struct node *self) {
  if (self != NULL) {
    self->height = 1;
    if (self->lhs != NULL) self->height = self->lhs->height + 1;
    if (self->rhs != NULL && self->rhs->height + 1 > self->height) {
      self->height = self->rhs->height + 1;
    }
  }
}
struct node *node_rot_right(struct node *self) {
  if (self == NULL) return NULL;
  if (self->lhs == NULL) return self;

  /* I really hope self is right,
   * it's gonna be hell to debug if it's not...
   */

  struct node *left = self->lhs;

  if (self->par != NULL && self == self->par->lhs) {
    self->par->lhs = left;
  } else if (self->par != NULL && self == self->par->rhs) {
    self->par->rhs = left;
  }
  left->par = self->par;
  self->par = left;

  self->lhs = left->rhs;
  left->rhs = self;
  if (self->lhs != NULL) self->lhs->par = self;

  node_update_height(self);
  node_update_height(left);

  return left;
}
struct node *node_rot_left(struct node *self) {
  if (self == NULL) return NULL;
  if (self->rhs == NULL) return self;

  /* I really hope self is right,
   * it's gonna be hell to debug if it's not...
   */

  struct node *right = self->rhs;

  if (self->par != NULL && self == self->par->lhs) {
    self->par->lhs = right;
  } else if (self->par != NULL && self == self->par->rhs) {
    self->par->rhs = right;
  }
  right->par = self->par;
  self->par = right;

  self->rhs = right->lhs;
  right->lhs = self;
  if (self->rhs != NULL) self->rhs->par = self;

  node_update_height(self);
  node_update_height(right);

  return right;
}
void node_balance(struct node *self) {
  while (self != NULL) {
    node_update_height(self);

    int bf = self->lhs == NULL ? 0 : self->lhs->height;
    bf -= self->rhs == NULL ? 0 : self->rhs->height;

    if (bf > 1) {
      // left-heavy
      self = node_rot_right(self);
    } else if (bf < -1) {
      // right-heavy
      self = node_rot_left(self);
    }

    self = self->par;
  }
}
void node_insert(struct node *self, int val) {
  if (val == self->val) return;
  else if (val < self->val) {
    if (self->lhs == NULL) {
      self->lhs = node_new(val, self);
      node_balance(self);
    }
    else node_insert(self->lhs, val);
  } else {
    if (self->rhs == NULL) {
      self->rhs = node_new(val, self);
      node_balance(self);
    } else node_insert(self->rhs, val);
  }
}
struct node *node_next(struct node *self) {
  if (self->rhs != NULL) {
    struct node *res = self->rhs;
    while (res->lhs != NULL) res = res->lhs;
    return res;
  } else {
    struct node *prev = self;
    struct node *res = self->par;
    while (res != NULL && res->rhs == prev) {
      prev = res;
      res = res->par;
    }
    return res;
  }
}
void node_print(struct node *self, FILE *file) {
  if (self == NULL) {
    fputc('*', file);
    return;
  }
  if (self->lhs == NULL && self->rhs == NULL) {
    fprintf(file, "%d", self->val);
    return;
  }
  fputc('<', file);
  node_print(self->lhs, file);
  fprintf(file, " %d ", self->val);
  node_print(self->rhs, file);
  fputc('>', file);
}
void node_printheight(struct node *self, FILE *file) {
  if (self == NULL) {
    fputc('0', file);
    return;
  }
  if (self->lhs == NULL && self->rhs == NULL) {
    fprintf(file, "%d", self->height);
    return;
  }
  fputc('<', file);
  node_printheight(self->lhs, file);
  fprintf(file, " %d ", self->height);
  node_printheight(self->rhs, file);
  fputc('>', file);
}
void node_printbf(struct node *self, FILE *file) {
  if (self == NULL) {
    fputc('0', file);
    return;
  }
  if (self->lhs == NULL && self->rhs == NULL) {
    fputc('0', file);
    return;
  }
  fputc('<', file);
  node_printbf(self->lhs, file);
  int bf = self->lhs == NULL ? 0 : self->lhs->height;
  bf -= self->rhs == NULL ? 0 : self->rhs->height;
  fprintf(file, " %d ", bf);
  node_printbf(self->rhs, file);
  fputc('>', file);
}

struct set {
  struct node *head;
  size_t size;
};

void set_init(struct set *self) {
  self->head = NULL;
  self->size = 0;
}
void set_deinit(struct set *self) {
  node_free(self->head);
  self->head = NULL;
  self->size = 0;
}
void set_add(struct set *self, int val) {
  if (self->head == NULL) {
    self->head = node_new(val, NULL);
  } else {
    node_insert(self->head, val);
    while (self->head->par != NULL) self->head = self->head->par;
  }
  ++self->size;
}
struct node *set_first(struct set *self) {
  struct node *res = self->head;
  while (res->lhs != NULL) res = res->lhs;
  return res;
}
void set_print(struct set *self, FILE *file) {
  fputc('{', file);
  node_print(self->head, file);
  fputc('}', file);
}
void set_printheight(struct set *self, FILE *file) {
  fputc('|', file);
  node_printheight(self->head, file);
  fputc('|', file);
}
void set_printbf(struct set *self, FILE *file) {
  fputc(':', file);
  node_printbf(self->head, file);
  fputc(':', file);
}

int main() {
  int n;
  struct set set;
  set_init(&set);

  puts("Please enter the numbers you like. cntl+D to finish");

  for (;;) {
    fputs("> ", stdout);
    scanf("%d", &n);

    if (feof(stdin)) {
      putchar('\n');
      break;
    }

    set_add(&set, n);
    #ifdef DEBUG
    set_print(&set, stdout);
    putchar('\n');
    set_printheight(&set, stdout);
    putchar('\n');
    set_printbf(&set, stdout);
    putchar('\n');
    #endif
  }

  fputs("You like the following numbers:", stdout);

  for (struct node *nd = set_first(&set); nd != NULL; nd = node_next(nd)) {
    printf(" %d", nd->val);
  }

  putchar('\n');
}