(TOJ) 245 / 逆序數對2

http://sprout.tw/oj/pro/245/

我的想法：
用treap維護好之前的數字，然後查詢 -q+aj ~ -p+aj 之間的數字XD

#include <iostream>
#include <stdio.h>
#include <ctime>
#include <cstdlib>
using namespace std;

struct Treap {
 Treap* lc;
 Treap* rc;
 int key,pri,size;
 Treap(int _key) {
  lc=rc=NULL;
  key=_key;
  pri=rand();
  size=1;
 }
};

int Size(Treap* t) {
 return t ? t->size : 0;
}

void pull(Treap* t) {
 t->size = 1 + Size(t->lc) + Size(t->rc);
}

Treap* merge(Treap* a,Treap* b) {
 if (!a||!b) return a?a:b;
 else if (a->pri > b->pri) {
  a->rc=merge(a->rc,b);
  pull(a);
  return a;
 }
 else {
  b->lc=merge(a,b->lc);
  pull(b);
  return b;
 }
}

void split(Treap* t,int k,Treap* &a,Treap* &b) {
 if (!t) a=b=NULL;
 else if (t->key <= k) {
  a=t;
  split(t->rc,k,a->rc,b);
  pull(a);
 }
 else {
  b=t;
  split(t->lc,k,a,b->lc);
  pull(b);
 }
}

//p<=aj-ai<=q
//-q+aj<=ai<=-p+aj

int query(Treap* root,int val,int p,int q) {
 Treap* tl;
 Treap* tr;
// cout<<-q+val-1<<" ~ "<<-p+val <<endl;
 split(root,-q+val-1,tl,root);
 split(root,-p+val,root,tr);
 if (root!=NULL) {
//  cout<<"root->key = "<<root->key<<endl;
  pull(root);
 }
 else {
//  cout<<"empty\n";
 }
 int ret=Size(root);
 root=merge(merge(tl,root),tr);
 return ret;
}

void dfs(Treap* t) {
 cout<<t->key<<' ';
 if (t->lc) dfs(t->lc);
 if (t->rc) dfs(t->rc);
}

int main () {
 srand(time(NULL));
 int T;
 scanf("%d",&T);
 while (T--) {
  int n,p,q;
  scanf("%d %d %d",&n,&p,&q);
  long long ans=0;
  Treap* root=NULL;
  for (int x=0;n>x;x++) {
//   cout<<"x="<<x<<endl;
   int i;
   scanf("%d",&i);
   ans+=query(root,i,p,q);
   //root=merge(root,new Treap(i)); ==> 錯的 
   Treap* tl;
   Treap* tr;
   split(root,i-1,tl,root);
   split(root,i,root,tr);
   root=merge(root,new Treap(i));
   root=merge(merge(tl,root),tr);
  }
  
  printf("%lld\n",ans);
 }
}