Cod sursa(job #2610681)

Utilizator FLORENTIN-GIULIANO.DUMITRUDumitru Florentin Giuliano FLORENTIN-GIULIANO.DUMITRU Data 5 mai 2020 13:46:46
Problema Hashuri Scor 70
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 7.62 kb
#include <fstream>
#include <iostream>
std::ifstream fin("hashuri.in");
std::ofstream fout("hashuri.out");

struct node{
    bool  color; /// true = black, false = Red
    long val;
    node *left, *right, *parent;
    node (long val, bool col, node * a, node * b, node * c) :val(val), color(col), left(a),right(b), parent(c){}
};

class RBT{
protected:
    node * root;
    node * stop;
    void LeftRotate(node* );
    void RightRotate(node* );
    void repara_add(node* );
    void swap(node* , node* );
    void repara_remove(node* );
public:
    RBT();
    void search(long& );
    void add(long&);
    void remove(long& );
};

RBT::RBT() {
    stop = new node(0,true,NULL,NULL,NULL);
    root = stop;
}

void RBT::search(long &val){
    node* x = root;
    while(x != stop && x -> val != val){
        if(val < x -> val)
            x = x -> left;
        else
            x = x -> right;
    }
    if(x != stop)
        fout << "1\n";
    else
        fout <<  "0\n";
}

void RBT::LeftRotate(node* x){
    node*y  = x -> right;
    x -> right = y -> left;
    if(y -> left != stop)
        y -> left -> parent = x;
    y -> parent = x -> parent;
    if(x -> parent== stop)
        root = y;
    else if(x == x -> parent -> left)
        x -> parent -> left = y;
    else
        x -> parent -> right = y;
    y -> left = x;
    x -> parent = y;
}

void RBT::RightRotate(node* x){
    node*y = x -> left;
    x -> left = y -> right;
    if(y -> right != stop)
        y -> right -> parent = x;
    y -> parent = x -> parent;
    if(x -> parent == stop)
        root = y;
    else if(x == x -> parent -> right)
        x -> parent-> right = y;
    else
        x -> parent -> left = y;
    y -> right = x;
    x -> parent = y;
}

void RBT::repara_add(node* curent){
    while (!curent -> parent -> color) {
        if (curent->parent == curent->parent->parent->left) {
            node *auxiliar = curent->parent->parent->right;
            if (!auxiliar->color) {
                curent->parent->color = true;
                auxiliar->color = true;
                curent->parent->parent->color = false;
                curent = curent->parent->parent;
            } else {
                if (curent == curent->parent->right) {
                    curent = curent->parent;
                    LeftRotate(curent);
                }
                curent->parent->color = true;
                curent->parent->parent->color = false;
                RightRotate(curent->parent->parent);
            }
        } else {
            node *auxiliar = curent->parent->parent->left;
            if (!auxiliar->color) {
                curent->parent->color = true;
                auxiliar->color = true;
                curent->parent->parent->color = false;
                curent = curent->parent->parent;
            } else {
                if (curent == curent->parent->left) {
                    curent = curent->parent;
                    RightRotate(curent);
                }
                curent->parent->color = true;
                curent->parent->parent->color = false;
                LeftRotate(curent->parent->parent);
            }
        }
    }
    root->color = true;
}

void RBT::add(long& val) {
    node * nou = new node(val,false,stop,stop,stop);
    node* y = NULL;
    node* x = this->root;
    while (x != stop) {
        y = x;
        if (nou->val < x->val) {
            x = x->left;
        } else {
            x = x->right;
        }
    }
    nou->parent = y;
    if (y == NULL) {
        nou->parent = stop;
        root = nou;
    } else if (nou->val < y->val) {
        y->left = nou;
    } else {
        y->right = nou;
    }
    if (nou->parent == stop){
        nou->color = true;
        return;
    }
    if (nou->parent->parent == NULL) {
        return;
    }
    repara_add(nou);
}

void RBT::swap(node* first, node* second){
    if(first -> parent == stop)
        root = second;
    else if(first == first -> parent -> left)
        first -> parent -> left = second;
    else
        first -> parent -> right = second;
    second -> parent = first -> parent;
}

void RBT::repara_remove(node* curent){
    while(curent != root && curent -> color) {
        if (curent == curent->parent->left) {
            node *auxiliar = curent->parent->right;
            if (!auxiliar->color) {
                auxiliar->color = true;
                curent->parent->color = false;
                LeftRotate(curent->parent);
                auxiliar = curent->parent->right;
            }
            if (auxiliar->left->color && auxiliar->right->color) {
                auxiliar->color = false;
                curent = curent->parent;
            } else {
                if (auxiliar->right->color) {
                    auxiliar->left->color = true;
                    auxiliar->color = false;
                    RightRotate(auxiliar);
                    auxiliar = curent->parent->right;
                }
                auxiliar->color = curent->parent->color;
                curent->parent->color = true;
                auxiliar->right->color = true;
                LeftRotate(curent->parent);
                curent = root;
            }
        } else {
            node *auxiliar = curent->parent->left;
            if (!auxiliar->color) {
                auxiliar->color = true;
                curent->parent->color = false;
                RightRotate(curent->parent);
                auxiliar = curent->parent->left;
            }
            if (auxiliar->left->color && auxiliar->right->color) {
                auxiliar->color = false;
                curent = curent->parent;
            } else {
                if (auxiliar->left->color) {
                    auxiliar->right->color = true;
                    auxiliar->color = false;
                    LeftRotate(auxiliar);
                    auxiliar = curent->parent->left;
                }
                auxiliar->color = curent->parent->color;
                curent->parent->color = true;
                auxiliar->left->color = true;
                RightRotate(curent->parent);
                curent = root;
            }
        }
    }
    curent->color = true;
}

void RBT::remove(long& val){
    node* z = root;
    while(z != stop && z -> val != val){
        if(val < z -> val)
            z = z -> left;
        else
            z = z -> right;
    }
    if(z != stop) {
        node* x, *y = z;
        bool color_originala = y -> color;
        if(z -> left == stop){
            x = z -> right;
            swap(z, z -> right);
        }
        else if(z -> right == stop){
            x = z -> left;
            swap(z, z -> left);
        }
        else{
            y = z -> right;
            while(y -> left != stop)
                y = y -> left;
            color_originala = y -> color;
            x = y -> right;
            if(y -> parent == z)
                x -> parent = y;
            else{
                swap(y, y -> right);
                y -> right = z -> right;
                y -> right -> parent = y;
            }
            swap(z, y);
            y -> left = z -> left;
            y -> left -> parent = y;
            y -> color = z -> color;
        }
        if(color_originala)
            repara_remove(x);
    }
}

int main() {
    RBT a;
    int n;
    fin >> n;
    while(n) {
        n--;
        int operatie;
        fin >> operatie;
        if (operatie == 1) {
            long val;
            fin >> val;
            a.add(val);
        } else if (operatie == 2) {
            long val;
            fin >> val;
            a.remove(val);
        } else if (operatie == 3) {
            long val;
            fin >> val;
            a.search(val);
        }
    }
    fin.close();
    fout.close();
    return 0;
}