Advertisement
cosenza987

Untitled

Mar 10th, 2024
573
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 5.71 KB | None | 0 0
  1. //Слава Україні, Героям слава
  2.  
  3. #include <bits/stdc++.h>
  4. #include <ext/pb_ds/assoc_container.hpp>
  5.  
  6. using namespace std;
  7. using namespace __gnu_pbds;
  8.  
  9. typedef tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update> indexed_set;
  10.  
  11. const int L = 17;
  12.  
  13. struct seg_tree {
  14.     int n;
  15.     vector<int> lz, st;
  16.     seg_tree(int n_ = 1) {
  17.         n = n_;
  18.         lz.resize(4 * n);
  19.         st.resize(4 * n);
  20.     }
  21.     void push(int p, int l, int r) {
  22.         if(lz[p]) {
  23.             st[p] += lz[p] * (r - l + 1);
  24.             if(l != r) {
  25.                 lz[2 * p] += lz[p];
  26.                 lz[2 * p + 1] += lz[p];
  27.             }
  28.             lz[p] = 0;
  29.         }
  30.     }
  31.     int query(int i, int j) {
  32.         return query(i, j, 1, 1, n);
  33.     }
  34.     int query(int i, int j, int p, int l, int r) {
  35.         push(p, l, r);
  36.         if(l > j or r < i) return 0;
  37.         if(l >= i and j >= r) return st[p];
  38.         return query(i, j, 2 * p, l, (l + r) >> 1) + query(i, j, 2 * p + 1, ((l + r) >> 1) + 1, r);
  39.     }
  40.     void update(int i, int j, int v) {
  41.         return update(i, j, v, 1, 1, n);
  42.     }
  43.     void update(int i, int j, int v, int p, int l, int r) {
  44.         push(p, l, r);
  45.         if(l > j or r < i) return;
  46.         if(l >= i and j >= r) {
  47.             lz[p] += v;
  48.             push(p, l, r);
  49.             return;
  50.         }
  51.         update(i, j, v, 2 * p, l, (l + r) >> 1);
  52.         update(i, j, v, 2 * p + 1, ((l + r) >> 1) + 1, r);
  53.         st[p] = st[2 * p] + st[2 * p + 1];
  54.     }
  55. };
  56.  
  57. int main() {
  58.     ios_base::sync_with_stdio(false);
  59.     cin.tie(nullptr);
  60.     int n, qq;
  61.     cin >> n >> qq;
  62.     vector<vector<int>> adj(n + 1), up(L, vector<int>(n + 1));
  63.     vector<vector<pair<int, int>>> fq(n + 1);
  64.     map<int, map<int, vector<pair<int, int>>>> sq;
  65.     vector<long long> q(qq), ans(qq);
  66.     vector<int> sz(n + 1), h(n + 1);
  67.     int root = -1;
  68.     for(int i = 1; i <= n; i++) {
  69.         int x;
  70.         cin >> x;
  71.         if(!x) {
  72.             root = i;
  73.             continue;
  74.         }
  75.         up[0][i] = x;
  76.         adj[x].push_back(i);
  77.     }
  78.     for(int i = 1; i < L; i++) {
  79.         for(int j = 1; j <= n; j++) {
  80.             up[i][j] = up[i - 1][up[i - 1][j]];
  81.         }
  82.     }
  83.     for(int i = 0; i < qq; i++) {
  84.         cin >> q[i];
  85.         q[i]--;
  86.         ans[i] += 1ll * (q[i] / n) * n * n;
  87.         fq[(q[i] / n) + 1].emplace_back(q[i] % n, i);
  88.     }
  89.     seg_tree seg(n);
  90.     function<void(int)> dfs_sz = [&](int u) {
  91.         sz[u] = 1;
  92.         for(auto e : adj[u]) {
  93.             h[e] = h[u] + 1;
  94.             dfs_sz(e);
  95.             sz[u] += sz[e];
  96.         }
  97.     };
  98.     dfs_sz(root);
  99.     seg.update(root, root, sz[root]);
  100.     function<int(int, int)> prv = [&](int a, int b) {
  101.         if(a == b) return a;
  102.         if(h[a] > h[b]) swap(a, b);
  103.         int dff = h[b] - h[a] - 1, cnt = 0;
  104.         while(dff) {
  105.             if(dff & 1) {
  106.                 b = up[cnt][b];
  107.             }
  108.             cnt++;
  109.             dff >>= 1;
  110.         }
  111.         if(up[0][b] == a) return b;
  112.         b = up[0][b];
  113.         for(int i = L - 1; i >= 0; i--) {
  114.             if(up[i][a] != up[i][b]) {
  115.                 a = up[i][a];
  116.                 b = up[i][b];
  117.             }
  118.         }
  119.         return a;
  120.     };
  121.     function<void(int)> dfsfq = [&](int u) {
  122.         for(auto [qqq, i] : fq[u]) {
  123.             int l = 1, r = n, res = -1;
  124.             while(l <= r) {
  125.                 int mid = (l + r) >> 1;
  126.                 if(seg.query(1, mid) > qqq) {
  127.                     res = mid;
  128.                     r = mid - 1;
  129.                 } else {
  130.                     l = mid + 1;
  131.                 }
  132.             }
  133.             ans[i] += 1ll * (res - 1) * n;
  134.             sq[res][prv(res, u)].emplace_back(qqq - seg.query(1, res - 1), i);
  135.         }
  136.         for(auto e : adj[u]) {
  137.             seg.update(e, e, sz[e]);
  138.             seg.update(u, u, -sz[e]);
  139.             dfsfq(e);
  140.             seg.update(e, e, -sz[e]);
  141.             seg.update(u, u, sz[e]);
  142.         }
  143.     };
  144.     dfsfq(root);
  145.     vector<vector<int>> v(n + 1);
  146.     indexed_set cur;
  147.     function<void(int, bool)> dfssq = [&](int u, bool keep) {
  148.         int mx = -1, bc = -1;
  149.         for(auto e : adj[u]) {
  150.             if(mx < sz[e]) {
  151.                 mx = sz[e];
  152.                 bc = e;
  153.             }
  154.         }
  155.         for(auto e : adj[u]) {
  156.             if(e != bc) {
  157.                 dfssq(e, false);
  158.             }
  159.         }
  160.         if(bc != -1) {
  161.             dfssq(bc, true);
  162.             swap(v[u], v[bc]);            
  163.         }
  164.         indexed_set tmp;
  165.         v[u].push_back(u);
  166.         cur.insert(u);
  167.         tmp.insert(u);
  168.         for(auto e : adj[u]) {
  169.             if(e != bc) {
  170.                 for(auto i : v[e]) {
  171.                     cur.insert(i);
  172.                     tmp.insert(i);
  173.                     v[u].push_back(i);
  174.                 }
  175.             }
  176.         }
  177.         for(auto [pos, id] : sq[u][u]) {
  178.             ans[id] += *(cur.find_by_order(pos)) - 1;
  179.         }
  180.         for(auto [pos, id] : sq[u][bc]) {
  181.             ans[id] += *(tmp.find_by_order(pos)) - 1;
  182.         }
  183.         for(auto e : adj[u]) {
  184.             if(e != bc) {
  185.                 for(auto i : v[e]) {
  186.                     cur.erase(i);
  187.                 }
  188.                 for(auto [pos, id] : sq[u][e]) {
  189.                     ans[id] += *(cur.find_by_order(pos)) - 1;
  190.                 }
  191.                 for(auto i : v[e]) {
  192.                     cur.insert(i);
  193.                 }
  194.             }
  195.         }
  196.         if(keep == false) {
  197.             for(auto e : v[u]) {
  198.                 cur.erase(e);
  199.             }
  200.         }
  201.     };
  202.     dfssq(root, false);
  203.     for(int i = 0; i < qq; i++) {
  204.         cout << ans[i] << "\n";
  205.     }
  206.     return 0;
  207. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement