PAT 1057. Stack (30)

http://www.patest.cn/contests/pat-a-practise/1057

treap 求第k大的数

#include<cstdio>
#include<stack>
#include<cstring>
#include<ctime>
#include<cstdlib>

using namespace std;

struct Node {
    Node *ch[2];
    int rank;
    int value;
    int cnt;
    int size;

    Node():rank(0), value(0), cnt(0), size(0) {}

    bool operator < (const Node& other)  const {
        return rank < other.rank;
    }

    int cmp(int x) const {
        if (x == value) {
            return -1;
        }
        return x < value? 0 : 1;
    }

    void maintain() {
        size = cnt;
        if (ch[0] != NULL) size += ch[0]->size;
        if (ch[1] != NULL) size += ch[1]->size;
    }
};


void rotate(Node* &o, int d) {
    Node* k = o->ch[d ^ 1];
    o->ch[d ^ 1] = k->ch[d];
    k->ch[d] = o;
    o->maintain();
    k->maintain();
    o = k;
}

void insert(Node* &o, int x) {
    if (o == NULL) {
        o = new Node();
        o->ch[0] = NULL;
        o->ch[1] = NULL;
        o->value = x;
        o->cnt = 1;
        o->size = 1;
        o->rank = rand();
#ifdef DEBUG
        printf("rank = %d\n", o->rank);
#endif
    } else {
        int d = o->cmp(x);
        if (d == -1) {
            o->cnt++;
        } else {
            insert(o->ch[d], x);
            if (o->rank < o->ch[d]->rank) {
                rotate(o, d ^ 1);
            }
        }
    }
    o->maintain();
}

void remove(Node* &o, int x) {
    int d = o->cmp(x);
    if (d == -1) {
        Node* u = o;
        if (o->cnt > 1) {
            o->cnt--;
        } else if (o->ch[0] != NULL && o->ch[1] != NULL) {
            int d2 = (o->ch[0]->rank > o->ch[1]->rank ? 1 : 0);
            rotate(o, d2);
            remove(o->ch[d2], x);
        } else {
            if (o->ch[0] == NULL) o = o->ch[1];
            else o = o->ch[0];
            delete u;
        }
    } else { 
        remove(o->ch[d], x);
    }
    if (o != NULL) o->maintain();
}

//for debug
int find(Node* o, int x) {
    while(o != NULL) {
        int d = o->cmp(x);
        if (d == -1) return 1;
        else o = o->ch[d];
    }
    return 0;
}

int kth(Node* o, int k) {
    if (o == NULL || k <= 0 || k > o->size) return 0;
    int s = o->ch[0] == NULL ? 0 : o->ch[0]->size;
    if (k > s && k <= s + o->cnt) {
        return o->value;
    } else if (k <= s) {
        return kth(o->ch[0], k);
    } else return kth(o->ch[1], k - s - o->cnt);
}

int mymax(int a, int b) {
    return a > b ? a : b;
}

//for debug
int max_depth(Node* root) {
    if (root == NULL) {
        return 0;
    }
    Node *p = root;
    int dep = 0;
    while(NULL != p) {
        p = p->ch[0];
        ++dep;
    }
    printf("left = %d    ", dep);
    p = root;
    dep = 0;
    while(NULL != p) {
        p = p->ch[1];
        ++dep;
    }
    printf("right = %d\n", dep);
}

Node* root = NULL;
stack<int> st;
int main() {
    root = NULL;
#ifdef DEBUG
    printf("rand_max = %d\n", RAND_MAX);
#endif
    freopen("input", "r",stdin);
    srand(time(NULL));
    char op[20];
    int n;
    scanf("%d", &n);
    for (int i = 0; i < n; ++i) {
        scanf("%s", op);
#ifdef DEBUG
        puts(op);
#endif
        if (strcmp(op, "Push") == 0) {
            int key;
            scanf("%d", &key);
            st.push(key);
            insert(root, key);
#ifdef DEBUG
            printf("find %d  %d\n", key, find(root, key));
#endif
        } else if (strcmp(op, "Pop") == 0) {
            if (st.empty()) {
                printf("Invalid\n");
            } else {
                int v = st.top();
                st.pop();
                remove(root, v);
                printf("%d\n", v);
            }
        } else if (strcmp(op, "PeekMedian") == 0) {
            if (st.empty()) {
                printf("Invalid\n");
            } else {
                int size = st.size();
                int k;
                if (size & 1) {
                    k = (size + 1) >> 1;
                } else k = size >> 1;
#ifdef DEBUG
                printf("k = %d\n", k);
#endif
                printf("%d\n", kth(root, k));
            }
        } else printf("Invalid\n");
    }
#ifdef DEBUG
    max_depth(root);
#endif
    return 0;
}

 

 

树状数组的解法：

#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<stack>

using namespace std;

const int MAXN = 100000 + 10;
int c[MAXN]; //binary indexed array

int lowbit(int x) {
    return x & (-x);
}

int sum(int x) {
    int sum = 0;
    while(x > 0) {
        sum += c[x];
        x -= lowbit(x);
    }
    return sum;
}

void add(int p, int x) {
    while(p < MAXN) {
        c[p] += x;
        p += lowbit(p);
    }
}

int kth(int k) { //find the kth number by binary search, remember find the leftmost 
    int low = 1, high = MAXN - 1;
    while(low < high) {
        int mid = low + ((high - low) >> 1);
        int s = sum(mid);
        if(s < k) {
            low = mid + 1;
        } else high = mid;
    }
    return low;
}

stack<int> st;
int main() {
    freopen("input", "r", stdin);
    memset(c, 0, sizeof(c));
    char op[20];
    int n;
    scanf("%d", &n);
    for (int kase = 0; kase < n; ++kase) {
        scanf("%s", op);
        if (op[1] == 'u') {
            int k;
            scanf("%d", &k);
            st.push(k);
            add(k, 1);
        } else if (op[1] == 'o') {
            if (st.empty()) {
                printf("Invalid\n");
            } else {
                int k = st.top();
                st.pop();
                add(k, -1);
                printf("%d\n", k);
            }
        } else {
            if (st.empty()) {
                printf("Invalid\n");
            } else {
                int sz = st.size();
                printf("%d\n", kth((sz + 1) >> 1));
            }
        }
    }
    return 0;
}