TreeList

As we know, array data structure has O(1) access time and O(n) insertion/deletion(kth element); linked list need O(n) for all operations, and O(1) insertion/deletion only when node pointer is given.

We also know that binary search tree and skip-list provide O(log n) for accessing, insertion and deletion.

Now what if we want a data structure, which behaves like a list, which can insert/delete element at kth position efficiently?

One of my friend suggest “blocked linked list”, which is a combination of array and linked list. Assume the input data size N is known, we can use a list of arrays to store data. List has length sqrt(N), each array also has length sqrt(N). Insertion and deletion has complexity O(sqrt(N)), and array is split when length exceed sqrt(N), and adjacent blocks are merged when length smaller than sqrt(N). This data structure is very popular in high school programming contest, where data size is known in advance. However, O(sqrt(N)) is not good enough: we want O(log n) for all operations.

Another method is to use binary search tree: binary search tree with “size” field can find “kth” element in O(log n) time. Elements have a floating point value as key for comparison. When inserting kth element, we first look for element k-1 and element k, and take their mean as key for new element. Thus new element is inserted into kth position. However, this method is not that elegant: size of the list depends on floating point precision.

Finally, comes our “TreeList” data structure. I don’t think I am the first guy who invent it, and I don’t know whether there is any other name for such data structure. The TreeList support 3 operations:

  • insert element into kth position
  • remove element in kth position
  • get kth element

The implementation is based on treap. Actually, any binary search tree can be used. The complexity for all 3 operations are O(log n). Source code is shown below.

When do we need this data structure? Follow my latest post

template<class T>
class TreeList{
  // list structure, support 3 operations:
  // -insert element into kth position
  // -remove element into kth position
  // -get kth element
  // implemented using treap
  struct node{
    T value;
    int pri,size;
    node *left,*right;
    node(){}
    node(const T& v,node* l,node* r){
      value=v;
      left=l;
      right=r;
      size=1;
      pri=rand();
    }
  }*root,*null;
  void resize(node* cur){
    if(cur!=null)cur->size=cur->left->size+cur->right->size+1;
  }
  void left_rotate(node* &cur){
    node *res=cur->left;
    cur->left=res->right;
    res->right=cur;
    resize(cur);
    cur=res;
    resize(cur);
  }
  void right_rotate(node* &cur){
    node *res=cur->right;
    cur->right=res->left;
    res->left=cur;
    resize(cur);
    cur=res;
    resize(cur);
  }
  void remove_node(node* &cur){
    if(cur->left == null && cur->right == null){
      cur=null;
    }else if(cur->left->pri < cur->right->pri){
      left_rotate(cur);
      remove_node(cur->right);
      resize(cur);
    }else{
      right_rotate(cur);
      remove_node(cur->left);
      resize(cur);
    }
  }

  void insert_to_kth_pos(const T& x,node* &cur,int k){
    if(cur == null && k==0){
      cur = new node(x,null,null);
    }else if(k <= cur->left->size){
      insert_to_kth_pos(x,cur->left,k);
      if(cur->pri > cur->left->pri)
        left_rotate(cur);
      resize(cur);
    }else if(k <= cur->size){
      insert_to_kth_pos(x,cur->right,k-cur->left->size-1);
      if(cur->pri > cur->right->pri)
        right_rotate(cur);
      resize(cur);
    }
  }
  void remove_kth_element(node* & cur,int k){
    if(k == cur->left->size){
      remove_node(cur);
    }else if(k < cur->left->size){
      remove_kth_element(cur->left,k);
      if(cur->pri > cur->left->pri)
        left_rotate(cur);
      resize(cur);
    }else if(k < cur->size){
      remove_kth_element(cur->right,k-cur->left->size-1);
      if(cur->pri > cur->right->pri)
        right_rotate(cur);
      resize(cur);
    }
  }
  void print(node* cur){
    if(cur == null)return;
    print(cur->left);
    cout<<cur->value<<" ";
    print(cur->right);
  }
  public:
  TreeList (){
    null=new node();
    null->pri = 0x7fffffff;
    null->size = 0;
    null->left=null->right=null;
    root=null;
  }
  void print(){
    print(root);
    cout<<endl;
  }
  void insert(const T& x,int k){
    insert_to_kth_pos(x,root,k);
  }
  void remove(int k){
    remove_kth_element(root,k);
  }
  T get(int k){
    node* cur=root;
    while(cur!=null){
      if(k < cur->left->size){
        cur=cur->left;
      }else if(k == cur->left->size){
        break;
      }else{
        k-=cur->left->size+1;
        cur=cur->right;
      }
    }
    return cur->value;
  }
};