fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. typedef long long ll;
  5.  
  6. struct Node {
  7.     ll start;          // first value of the segment
  8.     ll len;            // number of elements in the segment
  9.     ll seg_sum;        // sum of this segment: len*(2*start + len - 1)/2
  10.     ll size;           // total number of elements in the subtree
  11.     ll total_sum;      // sum of all elements in the subtree
  12.     Node *left, *right, *parent;
  13.  
  14.     Node(ll s, ll l) : start(s), len(l), left(nullptr), right(nullptr), parent(nullptr) {
  15.         seg_sum = l * (2*s + l - 1) / 2;
  16.         size = l;
  17.         total_sum = seg_sum;
  18.     }
  19. };
  20.  
  21. Node* root = nullptr;
  22. set<pair<ll, Node*>> segMap;   // maps start value to the node containing that segment
  23. ll next_A = 1;                 // next number to take from sequence A
  24.  
  25. // ---------- Splay tree utilities ----------
  26. void update(Node* x) {
  27.     if (!x) return;
  28.     x->size = x->len;
  29.     x->total_sum = x->seg_sum;
  30.     if (x->left) {
  31.         x->size += x->left->size;
  32.         x->total_sum += x->left->total_sum;
  33.     }
  34.     if (x->right) {
  35.         x->size += x->right->size;
  36.         x->total_sum += x->right->total_sum;
  37.     }
  38. }
  39.  
  40. void rotate(Node* x) {
  41.     Node* p = x->parent;
  42.     Node* g = p->parent;
  43.     if (p->left == x) {
  44.         p->left = x->right;
  45.         if (x->right) x->right->parent = p;
  46.         x->right = p;
  47.         p->parent = x;
  48.     } else {
  49.         p->right = x->left;
  50.         if (x->left) x->left->parent = p;
  51.         x->left = p;
  52.         p->parent = x;
  53.     }
  54.     update(p);
  55.     update(x);
  56.     x->parent = g;
  57.     if (g) {
  58.         if (g->left == p) g->left = x;
  59.         else g->right = x;
  60.     }
  61. }
  62.  
  63. // Splay x to the root of its tree, and update the tree root pointer.
  64. void splay(Node*& treeRoot, Node* x) {
  65.     if (!x) return;
  66.     while (x->parent) {
  67.         Node* p = x->parent;
  68.         Node* g = p->parent;
  69.         if (!g) {
  70.             rotate(x);
  71.         } else if ((g->left == p) == (p->left == x)) {
  72.             rotate(p);
  73.             rotate(x);
  74.         } else {
  75.             rotate(x);
  76.             rotate(x);
  77.         }
  78.     }
  79.     treeRoot = x;
  80. }
  81.  
  82. Node* find_max(Node* node) {
  83.     while (node->right) node = node->right;
  84.     return node;
  85. }
  86.  
  87. // Merge two trees: all elements in left are before those in right.
  88. Node* merge(Node*& left, Node*& right) {
  89.     if (!left) {
  90.         left = right;
  91.         return left;
  92.     }
  93.     if (!right) return left;
  94.     Node* max_left = find_max(left);
  95.     splay(left, max_left);   // now left is max_left, with no right child
  96.     left->right = right;
  97.     right->parent = left;
  98.     update(left);
  99.     return left;
  100. }
  101.  
  102. // ---------- Segment map management ----------
  103. void insert_into_segMap(Node* node) {
  104.     segMap.insert({node->start, node});
  105. }
  106.  
  107. void remove_from_segMap(Node* node) {
  108.     auto it = segMap.find({node->start, node});
  109.     if (it != segMap.end()) segMap.erase(it);
  110. }
  111.  
  112. // Find the segment that contains value v.
  113. Node* find_segment(ll v) {
  114.     auto it = segMap.upper_bound({v, nullptr});
  115.     if (it == segMap.begin()) return nullptr;
  116.     --it;
  117.     Node* node = it->second;
  118.     if (node->start <= v && v <= node->start + node->len - 1)
  119.         return node;
  120.     // Should not happen if the map is consistent
  121.     return nullptr;
  122. }
  123.  
  124. // ---------- Core operations ----------
  125.  
  126. // Split the tree rooted at 'treeRoot' into two trees:
  127. // left tree contains the first k elements, right tree the rest.
  128. pair<Node*, Node*> split_by_index(Node*& treeRoot, ll k) {
  129.     if (!treeRoot) return {nullptr, nullptr};
  130.     if (k == 0) {
  131.         Node* right = treeRoot;
  132.         right->parent = nullptr;
  133.         treeRoot = nullptr;
  134.         return {nullptr, right};
  135.     }
  136.     if (k == treeRoot->size) {
  137.         Node* left = treeRoot;
  138.         left->parent = nullptr;
  139.         treeRoot = nullptr;
  140.         return {left, nullptr};
  141.     }
  142.  
  143.     Node* cur = treeRoot;
  144.     ll kk = k;   // remaining elements to skip
  145.     while (true) {
  146.         ll left_size = (cur->left ? cur->left->size : 0);
  147.         if (kk <= left_size) {
  148.             cur = cur->left;
  149.         } else if (kk <= left_size + cur->len) {
  150.             break;
  151.         } else {
  152.             kk -= left_size + cur->len;
  153.             cur = cur->right;
  154.         }
  155.     }
  156.     // cur contains the k-th element
  157.     splay(treeRoot, cur);
  158.     // now treeRoot == cur
  159.     ll left_size = (treeRoot->left ? treeRoot->left->size : 0);
  160.     ll offset = kk - left_size;   // 1-based offset inside the node
  161.  
  162.     if (offset == treeRoot->len) {
  163.         // whole node goes to the left tree
  164.         Node* left = treeRoot->left;
  165.         Node* right = treeRoot->right;
  166.         treeRoot->left = treeRoot->right = nullptr;
  167.         update(treeRoot);
  168.         if (left) left->parent = nullptr;
  169.         if (right) right->parent = nullptr;
  170.         Node* L_tree = merge(left, treeRoot);
  171.         treeRoot = nullptr;
  172.         return {L_tree, right};
  173.     } else {
  174.         // split the node
  175.         Node* left_node = new Node(treeRoot->start, offset);
  176.         Node* right_node = new Node(treeRoot->start + offset, treeRoot->len - offset);
  177.         remove_from_segMap(treeRoot);
  178.         insert_into_segMap(left_node);
  179.         insert_into_segMap(right_node);
  180.  
  181.         Node* left_child = treeRoot->left;
  182.         Node* right_child = treeRoot->right;
  183.         delete treeRoot;
  184.  
  185.         Node* L_tree = merge(left_child, left_node);
  186.         Node* R_tree = merge(right_node, right_child);
  187.         if (L_tree) L_tree->parent = nullptr;
  188.         if (R_tree) R_tree->parent = nullptr;
  189.         treeRoot = nullptr;
  190.         return {L_tree, R_tree};
  191.     }
  192. }
  193.  
  194. // Delete an entire subtree and remove all its segments from the map.
  195. void delete_tree(Node* node) {
  196.     if (!node) return;
  197.     delete_tree(node->left);
  198.     delete_tree(node->right);
  199.     remove_from_segMap(node);
  200.     delete node;
  201. }
  202.  
  203. // Type 1: insert K numbers after value X
  204. void type1(ll X, ll K) {
  205.     Node* u = find_segment(X);
  206.     // Guaranteed that u exists
  207.     splay(root, u);
  208.     Node* L = root->left;
  209.     Node* R = root->right;
  210.     if (L) L->parent = nullptr;
  211.     if (R) R->parent = nullptr;
  212.     root->left = root->right = nullptr;
  213.  
  214.     ll d = X - root->start;
  215.     Node* left_node = nullptr;
  216.     Node* right_node = nullptr;
  217.     if (d + 1 < root->len) {
  218.         left_node = new Node(root->start, d + 1);
  219.         right_node = new Node(X + 1, root->len - (d + 1));
  220.     } else {
  221.         left_node = new Node(root->start, root->len);
  222.         right_node = nullptr;
  223.     }
  224.     remove_from_segMap(root);
  225.     insert_into_segMap(left_node);
  226.     if (right_node) insert_into_segMap(right_node);
  227.     delete root;
  228.  
  229.     Node* new_node = new Node(next_A, K);
  230.     next_A += K;
  231.     insert_into_segMap(new_node);
  232.  
  233.     Node* A = merge(L, left_node);
  234.     Node* B = merge(A, new_node);
  235.     if (right_node) {
  236.         Node* C = merge(B, right_node);
  237.         root = merge(C, R);
  238.     } else {
  239.         root = merge(B, R);
  240.     }
  241. }
  242.  
  243. // Type 2: delete H numbers immediately after value Y
  244. void type2(ll Y, ll H) {
  245.     Node* u = find_segment(Y);
  246.     splay(root, u);
  247.     Node* L = root->left;
  248.     Node* R = root->right;
  249.     if (L) L->parent = nullptr;
  250.     if (R) R->parent = nullptr;
  251.     root->left = root->right = nullptr;
  252.  
  253.     ll d = Y - root->start;
  254.     Node* left_node = nullptr;
  255.     Node* right_node = nullptr;
  256.     if (d + 1 < root->len) {
  257.         left_node = new Node(root->start, d + 1);
  258.         right_node = new Node(Y + 1, root->len - (d + 1));
  259.     } else {
  260.         left_node = new Node(root->start, root->len);
  261.         right_node = nullptr;
  262.     }
  263.     remove_from_segMap(root);
  264.     insert_into_segMap(left_node);
  265.     if (right_node) insert_into_segMap(right_node);
  266.     delete root;
  267.  
  268.     Node* A = merge(L, left_node);
  269.     Node* C = nullptr;
  270.     if (right_node) {
  271.         C = merge(right_node, R);
  272.     } else {
  273.         C = R;
  274.     }
  275.     if (!C) {
  276.         root = A;
  277.         return;
  278.     }
  279.     Node* C_root = C;
  280.     pair<Node*, Node*> split_result = split_by_index(C_root, H);
  281.     Node* P1 = split_result.first;
  282.     Node* P2 = split_result.second;
  283.     delete_tree(P1);
  284.     root = merge(A, P2);
  285. }
  286.  
  287. // Type 3: return sum of elements from index L to R (1-based)
  288. ll type3(ll L, ll R) {
  289.     pair<Node*, Node*> split1 = split_by_index(root, L - 1);
  290.     Node* left = split1.first;
  291.     Node* mid_right = split1.second;
  292.  
  293.     pair<Node*, Node*> split2 = split_by_index(mid_right, R - L + 1);
  294.     Node* mid = split2.first;
  295.     Node* right = split2.second;
  296.  
  297.     ll ans = mid->total_sum;
  298.     Node* temp = merge(left, mid);
  299.     root = merge(temp, right);
  300.     return ans;
  301. }
  302.  
  303. int main() {
  304.  
  305.     // Initial sequence: only one element 0
  306.     root = new Node(0, 1);
  307.     insert_into_segMap(root);
  308.  
  309.     int Q;
  310.     scanf("%d", &Q);
  311.     while (Q--) {
  312.         int op;
  313.         scanf("%d", &op);
  314.         if (op == 1) {
  315.             ll X, K;
  316.             scanf("%lld %lld", &X, &K);
  317.             type1(X, K);
  318.         } else if (op == 2) {
  319.             ll Y, H;
  320.             scanf("%lld %lld", &Y, &H);
  321.             type2(Y, H);
  322.         } else {
  323.             ll L, R;
  324.             scanf("%lld %lld", &L, &R);
  325.             printf("%lld\n", type3(L, R));
  326.         }
  327.     }
  328.  
  329.     // Cleanup (not strictly necessary)
  330.     delete_tree(root);
  331.     return 0;
  332. }
Success #stdin #stdout 0.01s 5320KB
comments (?)
stdin
copy 
6
1 0 3
1 1 2
3 3 5
2 4 2
1 4 4
3 2 7
stdout
copy 
6
24
https://ideone.com/azE5aA
language:
C++14 (gcc 8.3)
created:
3 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!