CPPSET - Tập hợp động

Tác giả: happyboy99x

Ngôn ngữ: C++

#include <iostream>
#include <cstdio>
using namespace std;

const int INF = (int) 1e9 + 7;

class SplayTree {
    struct Node {
        Node * child[2], * parent;
        int key;
    } * root, * nil;

    void setChild(Node * x, Node * y, int d) {
        x->child[d] = y;
        y->parent = x;
    }

    int getDirection(Node * x, Node * y) {
        return x->child[0] == y ? 0 : 1;
    }

    void rotate(Node * x, int d) {
        Node * y = x->child[d];
        Node * z = x->parent;
        setChild(x, y->child[d ^ 1], d);
        setChild(y, x, d ^ 1);
        setChild(z, y, getDirection(z, x));
    }

    Node * splay(Node * x) {
        if(x == nil) return nil;
        while(x->parent != nil) {
            Node * y = x->parent;
            Node * z = y->parent;
            int dy = getDirection(y, x);
            int dz = getDirection(z, y);
            if(z == nil) {
                rotate(y, dy);
            } else if(dy == dz) {
                rotate(z, dz);
                rotate(y, dy);
            } else {
                rotate(y, dy);
                rotate(z, dz);
            }
        }
        return x;
    }

    Node * leftMost(Node * x) const {
        if(x == nil) return nil;
        while(x->child[0] != nil)
            x = x->child[0];
        return x;
    }

    Node * rightMost(Node * x) const {
        if(x == nil) return nil;
        while(x->child[1] != nil)
            x = x->child[1];
        return x;
    }

    public:
    SplayTree() {
        root = nil = new Node();
        nil->child[0] = nil->child[1] = nil->parent = nil;
        nil->key = INF;
    }

    void insert(int key) {
        Node * curr = root;
        Node * last = nil;
        int dir = -1;
        while(curr != nil) {
            last = curr;
            if(key == curr->key) {
                root = splay(curr);
                return;
            } else if(key < curr->key) {
                curr = curr->child[0];
                dir = 0;
            } else {
                curr = curr->child[1];
                dir = 1;
            }
        }
        Node * u = new Node();
        u->key = key;
        u->child[0] = u->child[1] = nil;
        setChild(last, u, dir);
        root = splay(u);
    }

    void erase(int key) {
        if(!contain(key)) return;
        Node * left = root->child[0];
        Node * right = root->child[1];
        left->parent = right->parent = nil;
        delete root;
        if(left == nil) {
            root = right;
        } else {
            left = splay(rightMost(left));
            setChild(left, right, 1);
            root = left;
        }
    }

    bool isEmpty() const {
        return root == nil;
    }

    int minimum() {
        return (root = splay(leftMost(root)))->key;
    }

    int maximum() {
        return (root = splay(rightMost(root)))->key;
    }

    bool contain(int key) {
        Node * curr = root;
        while(curr != nil) {
            if(key == curr->key) {
                root = splay(curr);
                return true;
            }
            else if(key < curr->key) curr = curr->child[0];
            else curr = curr->child[1];
        }
        return false;
    }

    int greater(int x) {
        Node * curr = root;
        Node * res = nil;
        while(curr != nil) {
            if(x >= curr->key) {
                curr = curr->child[1];
            } else {
                if(res == nil || curr->key < res->key) res = curr;
                curr = curr->child[0];
            }
        }
        if(res != nil) root = splay(res);
        return res->key;
    }

    int smaller(int x) {
        Node * curr = root;
        Node * res = nil;
        while(curr != nil) {
            if(x <= curr->key) {
                curr = curr->child[0];
            } else {
                if(res == nil || curr->key > res->key) res = curr;
                curr = curr->child[1];
            }
        }
        if(res != nil) root = splay(res);
        return res->key;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    SplayTree tree;
    for(int type, x; cin >> type, type != 0; ) {
        if(type != 3 && type != 4) cin >> x;
        if(type == 1) {
            tree.insert(x);
        } else if(type == 2) {
            tree.erase(x);
        } else if(tree.isEmpty()) {
            puts("empty");
        } else if(type == 3) {
            printf("%d\n", tree.minimum());
        } else if(type == 4) {
            printf("%d\n", tree.maximum());
        } else if((type == 6 || type == 8) && tree.contain(x)) {
            printf("%d\n", x);
        } else {
            int res = type <= 6 ? tree.greater(x) : tree.smaller(x);
            if(res == INF) puts("no");
            else printf("%d\n", res);
        }
    }
    return 0;
}

Download