fork download
copy 
  1. #pragma GCC optimize("O3,unroll-loops")
  2. #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
  3.  
  4. #include <iostream>
  5. #include <vector>
  6. #include <algorithm>
  7.  
  8. using namespace std;
  9.  
  10. const int MOD = 1000000007;
  11.  
  12. // LCT Node
  13. struct Node {
  14.     int p, ch[2];
  15.     long long weight, sum_weight;
  16.     int vir_sz, total_sz;
  17. } tr[300005];
  18.  
  19. inline void pushup(int x) {
  20.     long long sw = tr[x].weight;
  21.     int ts = tr[x].vir_sz + 1;
  22.     if (tr[x].ch[0]) {
  23.         sw += tr[tr[x].ch[0]].sum_weight;
  24.         ts += tr[tr[x].ch[0]].total_sz;
  25.     }
  26.     if (tr[x].ch[1]) {
  27.         sw += tr[tr[x].ch[1]].sum_weight;
  28.         ts += tr[tr[x].ch[1]].total_sz;
  29.     }
  30.     tr[x].sum_weight = sw % MOD;
  31.     tr[x].total_sz = ts;
  32. }
  33.  
  34. inline bool is_root(int x) {
  35.     return tr[tr[x].p].ch[0] != x && tr[tr[x].p].ch[1] != x;
  36. }
  37.  
  38. inline void rotate(int x) {
  39.     int y = tr[x].p, z = tr[y].p;
  40.     int k = (tr[y].ch[1] == x);
  41.     if (!is_root(y)) tr[z].ch[tr[z].ch[1] == y] = x;
  42.     tr[x].p = z;
  43.     tr[y].ch[k] = tr[x].ch[k ^ 1];
  44.     if (tr[x].ch[k ^ 1]) tr[tr[x].ch[k ^ 1]].p = y;
  45.     tr[x].ch[k ^ 1] = y;
  46.     tr[y].p = x;
  47.     pushup(y);
  48.     pushup(x);
  49. }
  50.  
  51. inline void splay(int x) {
  52.     if (x == 0) return; // FIX: Prevent infinite loop on the dummy root
  53.     while (!is_root(x)) {
  54.         int y = tr[x].p, z = tr[y].p;
  55.         if (!is_root(y)) {
  56.             if ((tr[y].ch[1] == x) ^ (tr[z].ch[1] == y)) rotate(x);
  57.             else rotate(y);
  58.         }
  59.         rotate(x);
  60.     }
  61. }
  62.  
  63. inline void access(int x) {
  64.     if (x == 0) return; // FIX: Prevent virtual pathing on the dummy root
  65.     for (int y = 0; x; y = x, x = tr[x].p) {
  66.         splay(x);
  67.         if (tr[x].ch[1]) tr[x].vir_sz += tr[tr[x].ch[1]].total_sz;
  68.         if (y) tr[x].vir_sz -= tr[y].total_sz;
  69.         tr[x].ch[1] = y;
  70.         pushup(x);
  71.     }
  72. }
  73.  
  74. inline void cut(int x) {
  75.     access(x); splay(x);
  76.     tr[tr[x].ch[0]].p = 0;
  77.     tr[x].ch[0] = 0;
  78.     pushup(x);
  79. }
  80.  
  81. inline void link(int x, int y) {
  82.     access(x); splay(x);
  83.     access(y); splay(y);
  84.     tr[x].p = y;
  85.     tr[y].vir_sz += tr[x].total_sz;
  86.     pushup(y);
  87. }
  88.  
  89. struct Query {
  90.     int x, id;
  91.     bool operator<(const Query& other) const {
  92.         return x < other.x;
  93.     }
  94. };
  95.  
  96. int main() {
  97.     // Fast I/O
  98.     ios_base::sync_with_stdio(false);
  99.     cin.tie(NULL);
  100.  
  101.     int n, k;
  102.     if (!(cin >> n >> k)) return 0;
  103.  
  104.     vector<pair<int, int>> events;
  105.     for (int i = 1; i <= n; i++) {
  106.         int val;
  107.         cin >> val;
  108.         events.push_back({val, i});
  109.     }
  110.     // Sort array values ascending to know exactly when they deactivate
  111.     sort(events.begin(), events.end());
  112.  
  113.     int q;
  114.     cin >> q;
  115.     vector<Query> queries(q);
  116.     for (int i = 0; i < q; i++) {
  117.         cin >> queries[i].x;
  118.         queries[i].id = i;
  119.     }
  120.     // Sort queries ascending
  121.     sort(queries.begin(), queries.end());
  122.  
  123.     // 1. Precalculate initial sizes for x = 0 in O(N)
  124.     vector<int> initial_sz(n + 1, 1);
  125.     for (int i = n; i >= 1; i--) {
  126.         int p = max(0, i - k);
  127.         initial_sz[p] += initial_sz[i];
  128.     }
  129.  
  130.     // 2. Initialize the Link-Cut Tree optimally
  131.     for (int i = 1; i <= n; i++) {
  132.         tr[i].p = max(0, i - k);
  133.         tr[i].weight = i;
  134.         tr[i].sum_weight = i;
  135.         tr[i].total_sz = initial_sz[i];
  136.         tr[i].vir_sz = initial_sz[i] - 1; // Initially all children are virtual
  137.         tr[i].ch[0] = tr[i].ch[1] = 0;
  138.     }
  139.  
  140.     tr[0].p = 0;
  141.     tr[0].weight = 0;
  142.     tr[0].sum_weight = 0;
  143.     tr[0].total_sz = initial_sz[0];
  144.     tr[0].vir_sz = initial_sz[0] - 1;
  145.     tr[0].ch[0] = tr[0].ch[1] = 0;
  146.  
  147.     long long S = 0;
  148.     for (int i = 1; i <= n; i++) {
  149.         S = (S + 1LL * initial_sz[i] * i) % MOD;
  150.     }
  151.  
  152.     // 3. Process Queries Offline
  153.     vector<long long> ans(q);
  154.     int e_ptr = 0;
  155.  
  156.     for (int i = 0; i < q; i++) {
  157.         // While the active element flips to inactive
  158.         while (e_ptr < n && events[e_ptr].first <= queries[i].x) {
  159.             int idx = events[e_ptr].second;
  160.  
  161.             // Fetch properties before the cut
  162.             int old_p = max(0, idx - k);
  163.             access(old_p); splay(old_p);
  164.             long long dp_old = tr[old_p].sum_weight;
  165.  
  166.             // Disconnect from i-k
  167.             cut(idx);
  168.  
  169.             // Fetch exactly the isolated subtree size
  170.             access(idx); splay(idx);
  171.             long long sz_i = tr[idx].total_sz;
  172.  
  173.             // Nullify its weight
  174.             tr[idx].weight = 0;
  175.             pushup(idx);
  176.  
  177.             // Fetch properties of the new parent
  178.             int new_p = idx - 1;
  179.             access(new_p); splay(new_p);
  180.             long long dp_new = tr[new_p].sum_weight;
  181.  
  182.             // Connect to i-1
  183.             link(idx, new_p);
  184.  
  185.             // Mathematically deduce the exact delta on the total answer
  186.             long long delta = (dp_new - dp_old - idx) % MOD;
  187.             if (delta < 0) delta += MOD;
  188.  
  189.             S = (S + (sz_i % MOD) * delta) % MOD;
  190.  
  191.             e_ptr++;
  192.         }
  193.         ans[queries[i].id] = S;
  194.     }
  195.  
  196.     for (int i = 0; i < q; i++) {
  197.         cout << ans[i] << "\n";
  198.     }
  199.  
  200.     return 0;
  201. }
Success #stdin #stdout 0s 5324KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/l8MGWL
language:
C++14 (gcc 8.3)
created:
1 day 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!