Advertisement
_takumi

йц

Dec 22nd, 2022
813
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 4.22 KB | None | 0 0
  1. #include <iostream>
  2. #include <algorithm>
  3.  
  4. struct Node {
  5.     int size;
  6.     int value;
  7.     int height;
  8.     struct Node *left;
  9.     struct Node *right;
  10. };
  11.  
  12. class AVLTree {
  13. private:
  14.     Node *root_;
  15.  
  16.     void clear(Node *n) {
  17.         if (n == nullptr) {
  18.             return;
  19.         }
  20.         clear(n->left);
  21.         clear(n->right);
  22.         delete n;
  23.     }
  24.  
  25.     int height(Node *n) {
  26.         if (n == nullptr) {
  27.             return 0;
  28.         }
  29.         return n->height;
  30.     }
  31.  
  32.     int balance(Node *n) {
  33.         if (n == nullptr) {
  34.             return 0;
  35.         }
  36.         return height(n->left) - height(n->right);
  37.     }
  38.  
  39.     Node *insert(int val, Node *n) {
  40.         if (n == nullptr) {
  41.             n = new Node;
  42.             n->value = val;
  43.             n->left = nullptr;
  44.             n->right = nullptr;
  45.             n->height = 1;
  46.             n->size = 1;
  47.             return n;
  48.         }
  49.         if (val < n->value) {
  50.             n->left = insert(val, n->left);
  51.         } else if (val > n->value) {
  52.             n->right = insert(val, n->right);
  53.         } else {
  54.             return n;
  55.         }
  56.         update(n);
  57.         return rebalance(n);
  58.     }
  59.  
  60.     Node *erase(Node *n, int val) {
  61.         if (n == nullptr) {
  62.             return nullptr;
  63.         }
  64.         if (n->value > val) {
  65.             n->left = erase(n->left, val);
  66.         } else if (n->value < val) {
  67.             n->right = erase(n->right, val);
  68.         } else {
  69.             Node *left = n->left;
  70.             Node *right = n->right;
  71.             delete n;
  72.             if (left == nullptr) {
  73.                 return right;
  74.             } else if (right == nullptr) {
  75.                 return left;
  76.             }
  77.             Node *tmp = right;
  78.             while (tmp->left != nullptr) {
  79.                 tmp = tmp->left;
  80.             }
  81.             tmp->left = left;
  82.             update(tmp);
  83.             return rebalance(tmp);
  84.         }
  85.         update(n);
  86.         return rebalance(n);
  87.     }
  88.  
  89.     Node *rebalance(Node *n) {
  90.         int b = balance(n);
  91.         if (b > 1 && balance(n->left) < 0) {
  92.             n->left = leftRotate(n->left);
  93.             return rightRotate(n);
  94.         } else if (b < -1 && balance(n->right) > 0) {
  95.             n->right = rightRotate(n->right);
  96.             return leftRotate(n);
  97.         } else if (b > 1) {
  98.             return rightRotate(n);
  99.         } else if (b < -1) {
  100.             return leftRotate(n);
  101.         }
  102.         return n;
  103.     }
  104.  
  105.     Node *rightRotate(Node *n) {
  106.         Node *rotated = n->left;
  107.         Node *right = rotated->right;
  108.         rotated->right = n;
  109.         n->left = right;
  110.         update(n);
  111.         update(rotated);
  112.         return rotated;
  113.     }
  114.  
  115.     Node *leftRotate(Node *n) {
  116.         Node *rotated = n->right;
  117.         Node *left = rotated->left;
  118.         rotated->left = n;
  119.         n->right = left;
  120.         update(n);
  121.         update(rotated);
  122.         return rotated;
  123.     }
  124.  
  125.     void update(Node *n) {
  126.         n->height = std::max(height(n->left), height(n->right)) + 1;
  127.         n->size = size(n->left) + size(n->right) + 1;
  128.     }
  129.  
  130.     int size(Node *n) {
  131.         return n ? n->size : 0;
  132.     }
  133.  
  134.     int at(Node *n, int i) {
  135.         if (!n) {
  136.             return 0;
  137.         }
  138.         if (size(n->left) >= i) {
  139.             return at(n->left, i);
  140.         } else if (size(n->left) + 1 < i) {
  141.             return at(n->right, i - size(n->left) - 1);
  142.         } else {
  143.             return n->value;
  144.         }
  145.     }
  146.  
  147. public:
  148.     AVLTree() {
  149.         root_ = nullptr;
  150.     }
  151.  
  152.     void insert(int val) {
  153.         root_ = insert(val, root_);
  154.     }
  155.  
  156.     void erase(int value) {
  157.         root_ = erase(root_, value);
  158.     }
  159.  
  160.     int at(int index) {
  161.         return at(root_, index);
  162.     }
  163.  
  164.     ~AVLTree() {
  165.         clear(root_);
  166.     }
  167. };
  168.  
  169. int main() {
  170.     std::ios::sync_with_stdio(false);
  171.     std::cin.tie(nullptr);
  172.     AVLTree s;
  173.     int x, n;
  174.     char ch;
  175.     std::cin >> n;
  176.     for (int i = 0; i < n; ++i) {
  177.         std::cin >> ch >> x;
  178.         if (ch == '?') {
  179.             std::cout << s.at(x) << std::endl;
  180.         } else if (ch == '+') {
  181.             s.insert(x);
  182.         } else if (ch == '-') {
  183.             s.erase(x);
  184.         }
  185.     }
  186.     return 0;
  187. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement