leoanjos

Awesome Shawarma

Mar 8th, 2023
667
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 3.06 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. #define llong long long int
  6.  
  7. struct FenwickTree {
  8. private:
  9.     int n;
  10.     vector<int> tree;
  11.  
  12. public:
  13.     FenwickTree() {}
  14.     FenwickTree(int n) {
  15.         this->n = n + 1;
  16.         tree.assign(n + 1, 0);
  17.     }
  18.  
  19.     void add(int i, int v) {
  20.         if (!i) tree[0] += v;
  21.         else {
  22.             while (i < n) {
  23.                 tree[i] += v;
  24.                 i += LSO(i);
  25.             }
  26.         }
  27.     }
  28.  
  29.     int sum(int i, int j) {
  30.         return sum(j) - sum(i - 1);
  31.     }
  32.  
  33. private:
  34.     int LSO(int i) {
  35.         return i & -i;
  36.     }
  37.  
  38.     int sum(int i) {
  39.         if (i < 0) return 0;
  40.  
  41.         int ans = tree[0];
  42.         while (i) {
  43.             ans += tree[i];
  44.             i -= LSO(i);
  45.         }
  46.  
  47.         return ans;
  48.     }
  49. };
  50.  
  51. int N;
  52. vector<vector<int>> tree;
  53. vector<int> depth, nodes;
  54. FenwickTree cnt;
  55.  
  56. void get_depth_and_nodes(int node = 0, int parent = -1) {
  57.     nodes[node] = 1;
  58.     for (int child: tree[node])
  59.         if (child != parent) {
  60.             depth[child] = depth[node] + 1;
  61.             get_depth_and_nodes(child, node);
  62.             nodes[node] += nodes[child];
  63.         }
  64. }
  65.  
  66. void add(int node, int parent, int v) {
  67.     cnt.add(depth[node], v);
  68.     for (int child: tree[node])
  69.         if (child != parent)
  70.             add(child, node, v);
  71. }
  72.  
  73. llong count(int dist, int node, int parent) {
  74.     llong ans = cnt.sum(0, min(dist - depth[node], N));
  75.     for (int child: tree[node])
  76.         if (child != parent)
  77.             ans += count(dist, child, node);
  78.  
  79.     return ans;
  80. }
  81.  
  82. llong solve(int dist, int node = 0, int parent = -1, bool keep = false) {
  83.     int mx = 0, big = -1;
  84.     for (int child: tree[node])
  85.         if (child != parent && nodes[child] > mx) {
  86.             big = child;
  87.             mx = nodes[child];
  88.         }
  89.  
  90.     llong ans = 0LL;
  91.     for (int child: tree[node])
  92.         if (child != parent && child != big)
  93.             ans += solve(dist, child, node);
  94.  
  95.     if (big != -1) ans += solve(dist, big, node, true);
  96.  
  97.     for (int child: tree[node])
  98.         if (child != parent && child != big) {
  99.             ans += count(dist + 2 * depth[node], child, node);
  100.             add(child, node, 1);
  101.         }
  102.  
  103.     ans += cnt.sum(depth[node], min(depth[node] + dist, N));
  104.  
  105.     cnt.add(depth[node], 1);
  106.     if (!keep) add(node, parent, -1);
  107.  
  108.     return ans;
  109. }
  110.  
  111. int main() {
  112.     ios_base::sync_with_stdio(false);
  113.     cin.tie(NULL);
  114.  
  115.     freopen("awesome.in", "r", stdin);
  116.  
  117.     int T; cin >> T;
  118.     while (T--) {
  119.         int L, R;
  120.         cin >> N >> L >> R;
  121.  
  122.         tree.assign(N, vector<int>());
  123.         for (int i = 1; i < N; i++) {
  124.             int X, Y;
  125.             cin >> X >> Y;
  126.  
  127.             tree[X - 1].push_back(Y - 1);
  128.             tree[Y - 1].push_back(X - 1);
  129.         }
  130.  
  131.         nodes.resize(N);
  132.         depth.assign(N, 0);
  133.         get_depth_and_nodes();
  134.  
  135.         cnt = FenwickTree(N);
  136.         llong ans = solve(N - L - 1);
  137.         assert(cnt.sum(0, N) == 0);
  138.         ans -= solve(N - R - 2);
  139.  
  140.         cout << ans << "\n";
  141.     }
  142. }
Advertisement
Add Comment
Please, Sign In to add comment