poj 1442 treap

3532 ワード

//      ,     i   ,    i  
//   treap
#include <iostream>
#include <ctime>
#include <cstdlib>
#include <cstring>
#include <cstdio>

using namespace std;

#define MX 30010

int size, root;
struct Node
{
    int l, r, key, rand_fix;
    int countl, countr;
}node[MX];

//left  rotate
void l_rot(int &index)
{
    int y = node[index].r;
    node[index].r = node[y].l;
    node[index].countr = node[y].countl;
    node[y].l = index;
    node[y].countl = node[index].countl + node[index].countr + 1;
    index = y;
}
//right rotate
void r_rot(int &index)
{
    int y = node[index].l;
    node[index].l = node[y].r;
    node[index].countl = node[y].countr;
    node[y].r = index;
    node[y].countr = node[index].countl + node[index].countr + 1;
    index = y;
}
void insert(int &index, int nkey)
{
    if(index == -1)
    {
        index = ++size;
        node[index].l = node[index].r = -1;
        node[index].rand_fix = rand();
        node[index].key = nkey;
        node[index].countl = node[index].countr = 0;

        return ;
    }
    if(nkey < node[index].key)
    {
        node[index].countl++;
        insert(node[index].l, nkey);
        if(node[node[index].l].rand_fix > node[index].rand_fix)
            r_rot(index);
    }
    else
    {
        node[index].countr++;
        insert(node[index].r, nkey);
        if(node[node[index].r].rand_fix > node[index].rand_fix)
            l_rot(index);
    }
}
bool remove(int &index, int nkey)
{
    if(index == -1) return false;
    if(nkey < node[index].key)
        return remove(node[index].l, nkey);
    else if(nkey > node[index].key)
        return remove(node[index].r, nkey);
    if(node[index].l == -1 && node[index].r == -1)
        index = -1;
    else if(node[index].l == -1)
        index = node[index].r;
    else if(node[index].r == -1)
        index = node[index].l;
    else
    {
        if(node[node[index].l].rand_fix < node[node[index].r].rand_fix)
        {
            l_rot(index);
            remove(node[index].l, nkey);
        }
        else
        {
            r_rot(index);
            remove(node[index].r, nkey);
        }
    }
    return true;
}
int find(int index, int count)
{
   if(index == -1) return -1;
   if(node[index].countl + 1 == count ) return node[index].key;
   if(node[index].l == -1) return find(node[index].r, count-1);
   if(node[index].countl + 1 < count)   return find(node[index].r,count - (node[index].countl+1));
   
   return find(node[index].l, count);
}
        
void init()
{
    size = root = -1;
    srand(time(0));
}

void print(int index)
{
    printf("%d ", node[index].key);
    if(node[index].l != -1) printf("left: %d ",node[node[index].l].key);
    if(node[index].r != -1) printf("right:%d ",node[node[index].r].key);
    printf("count: %d  %d
", node[index].countl, node[index].countr); if(node[index].l != -1) print(node[index].l); if(node[index].r != -1) print(node[index].r); } int n, m; int s[MX]; int p, c, inc; int main() { //freopen("1.txt", "r", stdin); init(); c = 0; scanf("%d %d", &m, &n); for(int i = 0;i < m;i++) scanf("%d", s+i); for(int i = 0;i < n;i++) { scanf("%d", &p); while(c < p) insert(root,s[c++]); printf("%d
", find(root,++inc)); } // print(root); return 0; }