Randomized Binary Search Tree - arosh/arosh.github.com GitHub Wiki

struct T {
  int value;
  int size;
  T *left, *right;
};

typedef pair<T *, T *> P;

// だいたい256MBにおさまる予定
const int NODE_MAX = 10000000;
T node_buf[NODE_MAX];
int node_ptr = 0;

int node_size(T *x) {
  return (x == NULL) ? 0 : x->size;
}

T *update(T *x) {
  x->size = node_size(x->left) + 1 + node_size(x->right);
  return x;
}

T *create(int value, T *left, T *right) {
  if(node_ptr == NODE_MAX) abort();
  T *x = &node_buf[node_ptr++];
  x->value = value;
  x->left = left;
  x->right = right;
  return update(x);
}

T *merge(T *left, T *right, bool immutable) {
  if(left == NULL) return right;
  if(right == NULL) return left;
  int lsize = node_size(left);
  int rsize = node_size(right);
  if(rand() % (lsize + rsize) < lsize) {
    if(immutable) {
      return create(left->value, left->left, merge(left->right, right, immutable));
    }
    else {
      left->right = merge(left->right, right, immutable);
      return update(left);
    }
  }
  else {
    if(immutable) {
      return create(right->value, merge(left, right->left, immutable), right->right);
    }
    else {
      right->left = merge(left, right->left, immutable);
      return update(right);
    }
  }
}

P split(T *x, int pos, bool immutable) {
  if(x == NULL) return P(NULL, NULL);
  if(pos <= node_size(x->left)) {
    P p = split(x->left, pos, immutable);
    if(immutable) {
      return P(p.first, create(x->value, p.second, x->right));
    }
    else {
      x->left = p.second;
      return P(p.first, update(x));
    }
  }
  else {
    P p = split(x->right, pos - 1 - node_size(x->left), immutable);
    if(immutable) {
      return P(create(x->value, x->left, p.first), p.second);
    }
    else {
      x->right = p.first;
      return P(update(x), p.second);
    }
  }
}

int get(T *x, int pos) {
  if(node_size(x->left) == pos) return x->value;
  if(pos <= node_size(x->left)) {
    return get(x->left, pos);
  }
  else {
    return get(x->right, pos - 1 - node_size(x->left));
  }
}

T *insert(T *x, int pos, int value, bool immutable) {
  P p = split(x, pos, immutable);
  return merge(merge(p.first, create(value, NULL, NULL), immutable), p.second, immutable);
}

T *remove(T *x, int pos, bool immutable) {
  P p = split(x, pos, immutable);
  return merge(p.first, split(p.second, 1, immutable).second, immutable);
}