danielvitor23

Untitled

Nov 17th, 2021
990
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 1.79 KB | None | 0 0
  1. Code
  2. #include <bits/stdc++.h>
  3. using namespace std;
  4. const int MAX = 2e5+50;
  5.  
  6. int n, k1, k2;
  7. int64_t ans;
  8. vector<int> gr[MAX];
  9. int dad[MAX], sub[MAX];
  10. bool rem[MAX];
  11.  
  12. int b[MAX];
  13.  
  14. void init(int sz) {
  15.     for (int idx = 1; idx <= sz+1; ++idx)
  16.         b[idx] = 0;
  17. }
  18.  
  19. void add(int idx, int val, int sz) {
  20.     for (++idx; idx <= sz+1; idx += (idx & -idx)) {
  21.         b[idx] += val;
  22.     }
  23. }
  24.  
  25. int64_t sum(int idx, int sz) {
  26.     if (idx < 0) return 0;
  27.     idx = min(idx, sz);
  28.     int64_t res = 0;
  29.     for (++idx; idx > 0; idx -= (idx & -idx)) {
  30.         res += b[idx];
  31.     }
  32.     return res;
  33. }
  34.  
  35. int64_t query(int l, int r, int sz) {
  36.     return sum(r, sz)-sum(l-1, sz);
  37. }
  38.  
  39. int dfs(int u, int par) {
  40.     sub[u] = 1;
  41.     for (int to : gr[u]) if (to != par and !rem[to]) {
  42.         sub[u] += dfs(to, u);
  43.     }
  44.     return sub[u];
  45. }
  46.  
  47. int centroid(int u, int par, int sz) {
  48.     for (int to : gr[u])
  49.         if (to != par and !rem[to] and sub[to] > sz/2)
  50.             return centroid(to, u, sz);
  51.     return u;
  52. }
  53.  
  54. void getChildren(int u, int par, int d, vector<int> &v) {
  55.     v.push_back(d);
  56.     for (int to : gr[u]) {
  57.         if (!rem[to] and to != par) {
  58.             getChildren(to, u, d+1, v);
  59.         }
  60.     }
  61. }
  62.  
  63. int decomp(int u, int par) {
  64.     int sz = dfs(u, par);
  65.     int c = centroid(u, par, sz);
  66.     if (par == -1)
  67.         par = c;
  68.     rem[c] = true;
  69.     dad[c] = par;
  70.    
  71.     init(sz);
  72.     add(0, 1, sz);
  73.     for (int to : gr[c]) if (!rem[to]) {
  74.         vector<int> v;
  75.         getChildren(to, c, 1, v);
  76.         for (int x : v) {
  77.             ans += query(k1-x, k2-x, sz);
  78.         }
  79.         for (int x : v) {
  80.             add(x, 1, sz);
  81.         }
  82.     }
  83.  
  84.     for (int to : gr[c]) if (!rem[to])
  85.         decomp(to, c);
  86.     return c;
  87. }
  88.  
  89. int main() {
  90.     cin.tie(0)->sync_with_stdio(0);
  91.     cin >> n >> k1 >> k2;
  92.     for (int i = 0, a, b; i < n-1; ++i) {
  93.         cin >> a >> b, --a, --b;
  94.         gr[a].push_back(b);
  95.         gr[b].push_back(a);
  96.     }
  97.     decomp(0, -1);
  98.     cout << ans << '\n';
  99. }
Advertisement
Add Comment
Please, Sign In to add comment