【模板】普通平衡树

权值线段树30pts求调

#include <bits/stdc++.h>

namespace trees {
    class segment_tree {
    private:
        int how_many_nodes = 0;

        struct node {
            int key;
            node *left, *right;

            node() : key(1), left(nullptr), right(nullptr) {}
        } *root;

        void insert(int s, int t, int val, node *&p) {
            if (!p) {
                p = new node();
            } else {
                p->key++;
            }
            if (s == t) {
                how_many_nodes++;
                return;
            } else {
                int mid = (s + t) >> 1;
                if (val <= mid) {
                    insert(s, mid, val, p->left);
                } else {
                    insert(mid + 1, t, val, p->right);
                }
            }
        }

        void erase(int s, int t, int val, node *&p) {
            if (!p) {
                return;
            } else {
                if (s == t && t == val) {
                    how_many_nodes--;
                    p->key--;
                    if (p->key <= 0) {
                        delete p;
                        p = nullptr;
                    }
                    return;
                } else {
                    int mid = (s + t) >> 1;
                    if (val <= mid) {
                        erase(s, mid, val, p->left);
                    } else {
                        erase(mid + 1, t, val, p->right);
                    }
                    p->key = (p->left ? p->left->key : 0) + (p->right ? p->right->key : 0);
                    if (p->key <= 0) {
                        delete p;
                        p = nullptr;
                    }
                    return;
                }
            }
        }

        int queryval_rank(int s, int t, int val, node *p) {
            if (!p) {
                return 0;
            } else if (s == t) {
                if (s < val) {
                    return p->key;
                } else {
                    return 0;
                }
            } else {
                int mid = (s + t) >> 1;
                if (val <= mid) {
                    return queryval_rank(s, mid, val, p->left);
                } else {
                    return (p->left ? p->left->key : 0) + queryval_rank(mid + 1, t, val, p->right);
                }
            }
        }

        int querykth_element(int s, int t, int k, node *p) {
            if (!p) {
                return 123456789;
            } else if (s == t) {
                return s;
            } else {
                int left_size = (p->left ? p->left->key : 0), mid = (s + t) >> 1;
                if (k <= left_size) {
                    return querykth_element(s, mid, k, p->left);
                } else {
                    return querykth_element(mid + 1, t, k - left_size, p->right);
                }
            }
        }

        bool find(int s, int t, int val, node *p) {
            if (!p) {
                return false;
            } else if (s == t) {
                return true;
            } else {
                int mid = (s + t) >> 1;
                if (val <= mid) {
                    return find(s, mid, val, p->left);
                } else {
                    return find(mid + 1, t, val, p->right);
                }
            }
        }

    public:
        bool find(int val) {
            return find(-(1 << 24), 1 << 24, val, root);
        }

        void insert(int val) {
            insert(-(1 << 24), 1 << 24, val, root);
        }

        void erase(int val) {
            erase(-(1 << 24), 1 << 24, val, root);
        }

        int queryval_rank(int val) {
            return queryval_rank(-(1 << 24), 1 << 24, val, root) + 1;
        }

        int querykth_element(int k) {
            if (k < 1) {
                return -2147483647;
            } else if (k > how_many_nodes) {
                return 2147483647;
            } else {
                return querykth_element(-(1 << 24), 1 << 24, k, root);
            }
        }

        int pre(int x) {
            int rnk = queryval_rank(x);
//            if (find(x)) {
            return querykth_element(rnk - 1);
//            } else {
//            return querykth_element(rnk);
//            }
        }

        int nxt(int x) {
            int rnk = queryval_rank(x);
//            if (find(x)){
            return querykth_element(rnk + 1);
//            } else{
            return querykth_element(queryval_rank(x) + 1);
//            }
        }

        int size() {
            return how_many_nodes;
        }
    };
}
using namespace std;
trees::segment_tree tree;
int n;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n;
    for (int i = 0; i < n; ++i) {
        int op, x;
        cin >> op >> x;
        if (op == 1) {
            tree.insert(x);
        } else if (op == 2) {
            tree.erase(x);
        } else if (op == 3) {
            cout << tree.queryval_rank(x) << '\n';
        } else if (op == 4) {
            cout << tree.querykth_element(x) << '\n';
        } else if (op == 5) {
            cout << tree.pre(x) << '\n';
        } else {
            cout << tree.nxt(x) << '\n';
        }
    }
    return 0;
}