Data-Structure/Exercise/Next2/Q2.cpp

#include <iostream>
#include <algorithm>
#include <string>

using namespace std;

struct Node {
    int val, h;
    Node *l, *r;
    Node(int v) : val(v), h(1), l(nullptr), r(nullptr) {}
};

class AVLTree {
private:
    Node* root;
    string curOp;
    bool isRotateStep;
    bool isRotate;

    int getH(Node* n) { return n ? n->h : 0; }
    int getBal(Node* n) { return n ? getH(n->l) - getH(n->r) : 0; }
    void upH(Node* n) { if (n) n->h = max(getH(n->l), getH(n->r)) + 1; }

    Node* rotR(Node* y) {
        Node* x = y->l;
        y->l = x->r;
        x->r = y;
        upH(y); upH(x);
        return x;
    }

    Node* rotL(Node* x) {
        Node* y = x->r;
        x->r = y->l;
        y->l = x;
        upH(x); upH(y);
        return y;
    }

    Node* check(Node* n, int lv) {
        if (!n) return nullptr;
        upH(n);
        int b = getBal(n);

        if (b > 1) {
            isRotate = true;
            if (!isRotateStep) {
                cout << curOp << ": ";
                isRotateStep = true;
            }
            if (getBal(n->l) >= 0) { // LL
                cout << "Rebalance subtree rooted at node " << n->val << " on level " << lv << " with right rotation. ";
                return rotR(n);
            } else { // LR
                cout << "Rebalance subtree rooted at node " << n->val << " on level " << lv << " with left rotation and right rotation. ";
                n->l = rotL(n->l);
                return rotR(n);
            }
        }

        if (b < -1) {
            isRotate = true;
            if (!isRotateStep) {
                cout << curOp << ": ";
                isRotateStep = true;
            }
            if (getBal(n->r) <= 0) { // RR
                cout << "Rebalance subtree rooted at node " << n->val << " on level " << lv << " with left rotation. ";
                return rotL(n);
            } else { // RL
                cout << "Rebalance subtree rooted at node " << n->val << " on level " << lv << " with right rotation and left rotation. ";
                n->r = rotR(n->r);
                return rotL(n);
            }
        }
        return n;
    }

    Node* ins(Node* n, int k, int lv) {
        if (!n) return new Node(k);
        if (k < n->val) n->l = ins(n->l, k, lv + 1);
        else if (k > n->val) n->r = ins(n->r, k, lv + 1);
        else return n;
        return check(n, lv);
    }

    Node* del(Node* n, int k, int lv) {
        if (!n) return nullptr;
        if (k < n->val) n->l = del(n->l, k, lv + 1);
        else if (k > n->val) n->r = del(n->r, k, lv + 1);
        else {
            if (!n->l || !n->r) {
                Node* tmp = n->l ? n->l : n->r;
                delete n;
                return tmp;
            } else {
                Node* sub = n->r;
                while (sub->l) sub = sub->l;
                n->val = sub->val;
                n->r = del(n->r, sub->val, lv + 1);
            }
        }
        return check(n, lv);
    }

    void in(Node* n) {
        if (!n) return;
        in(n->l);
        cout << n->val << " ";
        in(n->r);
    }

    void pre(Node* n) {
        if (!n) return;
        cout << n->val << " ";
        pre(n->l);
        pre(n->r);
    }

    void clear(Node* n) {
        if (!n) return;
        clear(n->l);
        clear(n->r);
        delete n;
    }

public:
    AVLTree() : root(nullptr), isRotate(false) {}
    ~AVLTree() { clear(root); }

    void op(string type, int k) {
        curOp = type + " " + to_string(k);
        isRotateStep = false;
        if (type == "Insert") root = ins(root, k, 0);
        else root = del(root, k, 0);
        if (isRotateStep) cout << endl;
    }

    void show() {
        if (!root) return;
        if (isRotate) cout << endl;
        in(root); cout << endl << endl;
        pre(root); cout << endl;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    int T, cIdx = 1;
    while (cin >> T) {
        if (cIdx > 1) cout << endl;
        
        AVLTree tree;
        cout << "Case " << cIdx++ << ":" << endl;
        for (int i = 0; i < T; ++i) {
            string cmd; int k;
            cin >> cmd >> k;
            tree.op(cmd, k);
        }
        tree.show();
    }
    return 0;
}