gistfile1.cpp

#include <iostream>
#include <map>
#include <cstring>
#include <cstdio>
#include <utility>

using namespace std;

template<class T> struct splnode {
  splnode() : V(0), P(NULL) { C[0] = C[1] = NULL; fix(); }
  splnode(const T& V) : V(V), P(NULL) { C[0] = C[1] = NULL; fix(); }

  T V;
  splnode* P;
  splnode<T>* C[2];

  /* Extra augmented counters. */
  int count;

  /* Fix the parent pointers of our children.  Additionally if we have any
   * extra data we're tracking (e.g. sum of subtree elements) now is the time to
   * update them (all of the children will already be up to date). */
  void fix() {
    if(C[0]) C[0]->P = this;
    if(C[1]) C[1]->P = this;
    count = 1 + (C[0] ? C[0]->count : 0) + (C[1] ? C[1]->count : 0);
  }

  /* Simple zig step in the 'c' direction when we've reached the root. */
  void zig(int c) {
    splnode<T>* X = C[c];
    if(X->P = P) P->C[up()] = X;
    C[c] = X->C[1 - c];
    X->C[1 - c] = this;
    fix(); X->fix();
  }

  /* Zig zig in the 'c' direction both times. */
  void zigzig(int c) {
    splnode<T>* X = C[c];
    splnode<T>* Y = X->C[c];
    if(Y->P = P) P->C[up()] = Y;
    C[c] = X->C[1 - c];
    X->C[c] = Y->C[1 - c];
    Y->C[1 - c] = X;
    X->C[1 - c] = this;
    fix(); X->fix(); Y->fix();
  }

  /* Zig zag first in the 'c' direction then in the '1-c' direciton. */
  void zigzag(int c) {
    splnode<T>* X = C[c];
    splnode<T>* Y = X->C[1 - c];
    if(Y->P = P) P->C[up()] = Y;
    C[c] = Y->C[1 - c];
    X->C[1 - c] = Y->C[c];
    Y->C[1 - c] = this;
    Y->C[c] = X;
    fix(); X->fix(); Y->fix();
  }

  /* Return the direction of this relative its parent. */
  int up() {
    return !P ? -1 : (P->C[0] == this ? 0 : 1);
  }

  /* Splay this up to the root. */
  splnode<T>* splay() {
    for(; P; ) {
      int c1 = up();
      int c2 = P->up();
      if(c2 == -1) {
        P->zig(c1);
      } else if(c1 == c2) {
        P->P->zigzig(c1);
      } else {
        P->P->zigzag(c2);
      }
    }
    return this;
  }

  /* Return the max element of the subtree rooted at this. */
  splnode<T>* last() {
    return C[1] ? C[1]->last()  : splay();
  }

  /* Return the min element of the subtree rooted at this. */
  splnode<T>* first() {
    return C[0] ? C[0]->first() : splay();
  }
};

/* Find the smallest element greater than or equal to x. */
template<class T, class Compare>
splnode<T>* spl_lower_bound(splnode<T>*& rt, const T& x, Compare comp) {
  splnode<T>* v = rt;
  splnode<T>* res = NULL;
  while(v) {
    if(!comp(v->V, x)) {
      res = v;
      v = v->C[0];
    } else if(v->C[1]) {
      v = v->C[1];
    } else {
      rt = v->splay();
      break;
    }
  }
  return res ? (rt = res->splay()) : NULL;
}

template<class T>
splnode<T>* spl_lower_bound(splnode<T>*& rt, const T& x) {
  return spl_lower_bound(rt, x, less<T>());
}

/* Insert the tree 'x' before 'ins'.  If 'ins' is NULL insert 'x' at the end of
 * the tree. */
template<class T>
void spl_insert(splnode<T>*& rt, splnode<T>* ins, splnode<T>* x) {
  if(!rt) {
    rt = x;
  } else if(!ins) {
    rt = rt->last();
    rt->C[1] = x;
    rt->fix();
  } else {
    ins->splay();
    if(ins->C[0]) {
      x = x->first();
      x->C[0] = ins->C[0];
      ins->C[0] = x;
      x->fix(); ins->fix();
    } else {
      ins->C[0] = x;
      ins->fix();
    }
  }
}

/* Delete node x. */
template<class T>
void spl_erase(splnode<T>*& rt, splnode<T>* x) {
  x->splay();
  if(x->C[0] && x->C[1]) {
    rt = x->C[0];
    rt->P = NULL;
    rt = rt->last();
    rt->C[1] = x->C[1];
    rt->fix();
  } else if(x->C[0]) {
    rt = x->C[0];
  } else if(x->C[1]) {
    rt = x->C[1];
  }
  rt->P = NULL;
}

void dump(splnode<int>* x, int& idx) {
  if(!x) return;
  dump(x->C[0], idx);
  printf("%d: %d\n", idx++, x->V);
  dump(x->C[1], idx);
}

int main() {
  printf("Commands:\n"
         "  insert X\n"
         "  delete X\n"
         "  index X\n"
         "  display\n");

  char cmd[16];
  splnode<int>* rt = NULL;
  while(scanf("%10s", cmd) != EOF) {
    if(!strcmp(cmd, "insert")) {
      int x; scanf("%d", &x);
      splnode<int>* ins = spl_lower_bound(rt, x);
      spl_insert(rt, ins, new splnode<int>(x));
    } else if(!strcmp(cmd, "delete")) {
      int x; scanf("%d", &x);
      splnode<int>* v = spl_lower_bound(rt, x);
      if(v && v->V == x) {
        spl_erase(rt, v);
      } else {
        printf("element does not exist\n");
      }
    } else if(!strcmp(cmd, "index")) {
      int x; scanf("%d", &x);
      splnode<int>* v = spl_lower_bound(rt, x);
      if(v && v->V == x) {
        printf("%d\n", v->C[0] ? v->C[0]->count : 0);
      } else {
        printf("element does not exist\n");
      }
    } else if(!strcmp(cmd, "display")) {
      int idx = 0;
      dump(rt, idx);
    } else {
      printf("Unknown command\n");
    }
  }
  return 0;
}