fork download 
  1. #include <iostream>
  2. #include <vector>
  3. #include <tuple>
  4. #include <algorithm>
  5. #include <cassert>
  6.  
  7. const int N = 120005;
  8. const int C = 1000000;
  9.  
  10. inline int read_int()
  11. {
  12.     int x = 0;
  13.     char c;
  14.     do {
  15.         c = getchar_unlocked();
  16.     } while (c < '0');
  17.     do {
  18.         x = (x << 1) + (x << 3) + (c - '0');
  19.         c = getchar_unlocked();
  20.     } while (c >= '0');
  21.     return x;
  22. }
  23.  
  24. struct EdgeInterval
  25. {
  26.     int edge;
  27.     int s, e;
  28.  
  29.     EdgeInterval() = default;
  30.     EdgeInterval(const int _edge, const int _s, const int _e) : edge(_edge), s(_s), e(_e) {
  31.     }
  32. };
  33.  
  34. std::vector<EdgeInterval> intvs[C];
  35. std::vector<std::pair<int, int>> edge_queries[C]; // (time, edge)
  36. std::tuple<int, int, int> queries[N];
  37. std::pair<int, int> tree_edge[N - 1];
  38. std::pair<int, int> last_update[N - 1]; // (time, color)
  39.  
  40. long long query_answers[C];
  41. std::vector<int> color_updates[N + 1];
  42. long long sum_tree[2 * C];
  43. int idx[C];
  44.  
  45. // for solve_color
  46. std::vector<int> tree[4 * 4 * N];
  47. std::vector<std::pair<int, int>> normalized_queries[4 * N];
  48.  
  49. std::vector<int> deltas[C];
  50. long long *fvalue_at[C];
  51.  
  52. long long f(int x) 
  53. {
  54.     return static_cast<long long>(x) * (x - 1) / 2;
  55. }
  56.  
  57. struct DisjointSets
  58. {
  59.     int parent[N];
  60.     int weight[N];
  61.     std::vector<std::pair<int, int>> stk;
  62.  
  63.     long long fvalue;
  64.  
  65.     DisjointSets()
  66.     {
  67.         for (int i = 0; i < N; ++ i) {
  68.             parent[i] = i;
  69.             weight[i] = 1;
  70.         }
  71.         fvalue = 0;
  72.     }
  73.  
  74.     int find_root(int x) const {
  75.         return (parent[x] == x) ? x : find_root(parent[x]);    
  76.     }
  77.  
  78.     void unite(int x, int y) 
  79.     {
  80.         x = find_root(x);
  81.         y = find_root(y);
  82.         if (x == y) {
  83.             return;
  84.         }
  85.  
  86.         if (y > x) {
  87.             std::swap(x, y);
  88.         }
  89.         stk.emplace_back(x, y);
  90.         fvalue -= f(weight[x]);
  91.         fvalue -= f(weight[y]);
  92.  
  93.         weight[x] += weight[y];
  94.         parent[y] = x;
  95.         fvalue += f(weight[x]);
  96.     }
  97.  
  98.     int get_size(int node) const
  99.     {
  100.         return weight[find_root(node)];   
  101.     }
  102.  
  103.     int get_snapshot_idx() const 
  104.     {
  105.         return static_cast<int>(stk.size());
  106.     }
  107.  
  108.     void undo_at(const int snapshot) 
  109.     {
  110.         while (static_cast<int>(stk.size()) > snapshot) {
  111.             int x, y;
  112.             x = stk.back().first;
  113.             y = stk.back().second;
  114.  
  115.             fvalue -= f(weight[x]);
  116.             parent[y] = y;
  117.             weight[x] -= weight[y];
  118.             fvalue += f(weight[x]);
  119.             fvalue += f(weight[y]);
  120.  
  121.             stk.pop_back();
  122.         }
  123.     }
  124. } DS;
  125.  
  126. void update_tree(int node, int l, int r, int ql, int qr, int edgeidx)
  127. {
  128.     if (ql <= l && r <= qr) {
  129.         tree[node].push_back(edgeidx);
  130.     } else {
  131.         int mid = (l + r) >> 1;
  132.         if (ql <= mid) {
  133.             update_tree(2 * node + 1, l, mid, ql, qr, edgeidx);
  134.         }
  135.         if (mid < qr) {
  136.             update_tree(2 * node + 2, mid + 1, r, ql, qr, edgeidx);
  137.         }
  138.     }
  139. }
  140.  
  141. void dfs(const int color, int node, int l, int r) 
  142. {
  143.     const int snapshot = DS.get_snapshot_idx();
  144.     for (auto&& edgeidx : tree[node]) {
  145.         DS.unite(tree_edge[edgeidx].first, tree_edge[edgeidx].second);
  146.     }
  147.     tree[node].clear();
  148.  
  149.     if (l == r) {
  150.         fvalue_at[color][l] = DS.fvalue;
  151.         for (auto&& edge_query : normalized_queries[l]) {
  152.             query_answers[edge_query.first] = f(DS.get_size(tree_edge[edge_query.second].first));
  153.         }
  154.         normalized_queries[l].clear();
  155.     } else {
  156.         int mid = (l + r) >> 1;
  157.         dfs(color, 2 * node + 1, l, mid);
  158.         dfs(color, 2 * node + 2, mid + 1, r);
  159.     }
  160.  
  161.     DS.undo_at(snapshot);
  162. }
  163.  
  164. void solve_color(int color)
  165. {
  166.     for (auto&& x : intvs[color]) {
  167.         deltas[color].push_back(x.s);
  168.         deltas[color].push_back(x.e);
  169.         deltas[color].push_back(x.e + 1);
  170.     }
  171.     for (auto&& x : edge_queries[color]) {
  172.         deltas[color].push_back(x.first);
  173.     }
  174.     std::sort(deltas[color].begin(), deltas[color].end());
  175.     deltas[color].erase(std::unique(deltas[color].begin(), deltas[color].end()), deltas[color].end());
  176.  
  177.     int n = static_cast<int>(deltas[color].size());
  178.     fvalue_at[color] = new long long[n];
  179.     for (auto&& x : intvs[color]) {
  180.         int a = std::lower_bound(deltas[color].begin(), deltas[color].end(), x.s) - deltas[color].begin();
  181.         int b = std::lower_bound(deltas[color].begin(), deltas[color].end(), x.e) - deltas[color].begin();
  182.         update_tree(0, 0, n - 1, a, b, x.edge);
  183.     }
  184.     for (auto&& x : edge_queries[color]) {
  185.         normalized_queries[std::lower_bound(deltas[color].begin(), deltas[color].end(), x.first) - deltas[color].begin()].push_back(x);
  186.     }
  187.  
  188.     dfs(color, 0, 0, n - 1);
  189. }
  190.  
  191. int main()
  192. {
  193.     int n = read_int();
  194.     for (int i = 0; i < n - 1; ++ i) {
  195.         int x = read_int() - 1, y = read_int() - 1, color = read_int() - 1;
  196.         tree_edge[i] = std::make_pair(x, y);
  197.         last_update[i] = std::make_pair(-1, color);
  198.     }
  199.  
  200.     int q = read_int();
  201.     for (int i = 0; i < q; ++ i) {
  202.         int qtype = read_int();
  203.         if (qtype == 1) {
  204.             int edgeidx = read_int() - 1, color = read_int() - 1;
  205.             if (last_update[edgeidx].second != color) {
  206.                 intvs[last_update[edgeidx].second].emplace_back(edgeidx, last_update[edgeidx].first, i - 1);
  207.                 last_update[edgeidx] = std::make_pair(i, color);
  208.             }
  209.             queries[i] = std::make_tuple(-1, 0, 0);
  210.         } else {
  211.             if (qtype == 2) {
  212.                 int x = read_int() - 1, y = read_int() - 1;
  213.                 queries[i] = std::make_tuple(0, x, y);
  214.             } else {
  215.                 int edgeidx = read_int() - 1;
  216.                 queries[i] = std::make_tuple(1, edgeidx, -1);
  217.                 edge_queries[last_update[edgeidx].second].emplace_back(i, edgeidx);   
  218.             }
  219.         }
  220.     }
  221.  
  222.     for (int i = 0; i < n - 1; ++ i) {
  223.         intvs[last_update[i].second].emplace_back(i, last_update[i].first, q);
  224.     }
  225.     for (int i = 0; i < q; ++ i) {
  226.         query_answers[i] = -1;
  227.     }
  228.     for (int i = 0; i < C; ++ i) {
  229.         if (!intvs[i].empty()) {
  230.             //deltas[i].push_back(q);
  231.             solve_color(i);
  232.             for (auto&& x : deltas[i]) {
  233.                 if (x != -1) {
  234.                     color_updates[x].push_back(i);
  235.                 } else {
  236.                     sum_tree[i + C] = fvalue_at[i][0];
  237.                     idx[i] = 1;
  238.                 }
  239.             }
  240.         }
  241.     }
  242.  
  243.     for (int i = C - 1; i > 0; -- i) {
  244.         sum_tree[i] = sum_tree[i << 1] + sum_tree[i << 1 | 1];
  245.     }
  246.  
  247.     for (int i = 0; i < q; ++ i) {
  248.         for (auto&& color : color_updates[i]) {
  249.             sum_tree[color + C] = fvalue_at[color][idx[color]];
  250.             for (int node = color + C; node > 0; node >>= 1) {
  251.                 sum_tree[node >> 1] = sum_tree[node ^ 1] + sum_tree[node];
  252.             }
  253.  
  254.             idx[color]++;
  255.         }
  256.  
  257.         if (std::get<0>(queries[i]) == 0) {
  258.             query_answers[i] = 0ULL;
  259.             for (int l = C + std::get<1>(queries[i]), r = 1 + C + std::get<2>(queries[i]); l < r; l >>= 1, r >>= 1) {
  260.                 if (l & 1) {
  261.                     query_answers[i] += sum_tree[l++];
  262.                 }
  263.                 if (r & 1) {
  264.                     query_answers[i] += sum_tree[--r];
  265.                 }
  266.             }
  267.         }
  268.     }
  269.  
  270.     for (int i = 0; i < q; ++ i) {
  271.         if (query_answers[i] != -1) {
  272.             printf("%lld\n", query_answers[i]);
  273.         }
  274.     }
  275. }
Time limit exceeded #stdin #stdout 5s 184000KB
comments (?)
stdin 
Standard input is empty
stdout 
Standard output is empty
https://ideone.com/RWMxvT
language:
C++14 (gcc 8.3)
created:
7 years 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!