fork download 
  1. // ~~ icebear ~~
  2. #include <bits/stdc++.h>
  3. using namespace std;
  4. #define int ll
  5. typedef long long ll;
  6. typedef pair<int, int> ii;
  7. typedef pair<int, ii> iii;
  8.  
  9. template<class T>
  10.     bool minimize(T &a, const T &b) {
  11.         if (a > b) return a = b, true;
  12.         return false;
  13.     }
  14.  
  15. template<class T>
  16.     bool maximize(T &a, const T &b) {
  17.         if (a < b) return a = b, true;
  18.         return false;
  19.     }
  20.  
  21. #define FOR(i,a,b) for(int i=(a); i<=(b); ++i)
  22. #define FORR(i,a,b) for(int i=(a); i>=(b); --i)
  23. #define REP(i, n) for(int i=0; i<(n); ++i)
  24. #define RED(i, n) for(int i=(n)-1; i>=0; --i)
  25. #define MASK(i) (1LL << (i))
  26. #define BIT(S, i) (((S) >> (i)) & 1)
  27. #define mp make_pair
  28. #define pb push_back
  29. #define fi first
  30. #define se second
  31. #define all(x) x.begin(), x.end()
  32. #define task "xc"
  33.  
  34. const int MOD = 1e9 + 7;
  35. const int inf = 1e9 + 27092008;
  36. const ll INF = 1e18 + 27092008;
  37. const int N = 2e5 + 5;
  38. int n, k;
  39. vector<ii> G[N];
  40. int tour[N << 1], minHigh[N << 1][20], timer, tin[N];
  41. ll high[N];
  42.  
  43. struct Edge {
  44.     int u, v, c;
  45. } e[N];
  46.  
  47. void dfs(int u, int par) {
  48.     tin[u] = ++timer;
  49.     tour[timer] = u;
  50.     for(ii x : G[u]) {
  51.         int v, w; tie(v, w) = x;
  52.         if (v == par) continue;
  53.         high[v] = high[u] + w;
  54.         dfs(v, u);
  55.         tour[++timer] = u;
  56.     }
  57. }
  58.  
  59. #define MIN_HIGH(x, y) (high[x] < high[y] ? x : y)
  60. int LCA(int u, int v) {
  61.     int l = tin[u];
  62.     int r = tin[v];
  63.     if (l > r) swap(l, r);
  64.     int k = __lg(r - l + 1);
  65.     return MIN_HIGH(minHigh[l][k], minHigh[r - MASK(k) + 1][k]);
  66. }
  67.  
  68. ll dist(int u, int v) {
  69.     int p = LCA(u, v);
  70.     return high[u] + high[v] - 2 * high[p];
  71. }
  72.  
  73. struct DSU {
  74.     vector<int> lab, l_dia, r_dia;
  75.     vector<ll> diameter;
  76.     vector<array<int, 3>> history;
  77.     ll mx_dia;
  78.     DSU(int n = 0): lab(n + 5, -1), l_dia(n + 5, 0), r_dia(n + 5, 0), history(), diameter(n + 5, 0), mx_dia(0) {
  79.         FOR(i, 1, n) l_dia[i] = r_dia[i] = i;
  80.     }
  81.  
  82.     int root(int v) {
  83.         return (lab[v] < 0 ? v : root(lab[v]));
  84.     }
  85.  
  86.     bool unite(int u, int v) {
  87.         u = root(u);
  88.         v = root(v);
  89.         if (u == v) return false;
  90.         if (lab[u] > lab[v]) swap(u, v);
  91.         int l = l_dia[u], r = r_dia[u];
  92.         if (dist(l_dia[u], l_dia[v]) > dist(l, r)) l = l_dia[u], r = l_dia[v];
  93.         if (dist(l_dia[u], r_dia[v]) > dist(l, r)) l = l_dia[u], r = r_dia[v];
  94.         if (dist(r_dia[u], l_dia[v]) > dist(l, r)) l = r_dia[u], r = l_dia[v];
  95.         if (dist(r_dia[u], r_dia[v]) > dist(l, r)) l = r_dia[u], r = r_dia[v];
  96.         if (dist(l_dia[v], r_dia[v]) > dist(l, r)) l = l_dia[v], r = r_dia[v];
  97.  
  98.         history.pb({0, u, lab[u]});
  99.         history.pb({0, v, lab[v]});
  100.         history.pb({1, u, l_dia[u]});
  101.         history.pb({2, u, r_dia[u]});
  102.         history.pb({3, u, diameter[u]});
  103.         history.pb({4, -1, mx_dia});
  104.         lab[u] += lab[v];
  105.         lab[v] = u;
  106.         if (maximize(diameter[u], dist(l, r))) l_dia[u] = l, r_dia[u] = r;
  107.         maximize(mx_dia, diameter[u]);
  108.         return true;
  109.     }
  110.  
  111.     void snap_shot() {
  112.         history.pb({-1, -1, -1});
  113.     }
  114.  
  115.     void roll_back() {
  116.         while(!history.empty()) {
  117.             auto T = history.back(); history.pop_back();
  118.             if (T[0] == -1) break;
  119.             if (T[0] == 0) lab[T[1]] = T[2];
  120.             if (T[0] == 1) l_dia[T[1]] = T[2];
  121.             if (T[0] == 2) r_dia[T[1]] = T[2];
  122.             if (T[0] == 3) diameter[T[1]] = T[2];
  123.             if (T[0] == 4) mx_dia = T[2];
  124.         }
  125.     }
  126. } dsu;
  127. int ans[N];
  128.  
  129. void DnC(int qryL, int qryR, int ansL, int ansR) {
  130.     if (qryL > qryR) return;
  131.     int qryM = (qryL + qryR) >> 1, ansM = max(ansL, qryM + 1);
  132.     dsu.snap_shot();
  133.     FOR(i, qryL, qryM) dsu.unite(e[i].u, e[i].v);
  134.     if (dsu.mx_dia > k) {
  135.         dsu.roll_back();
  136.         DnC(qryL, qryM - 1, ansL, ansR);
  137.         return;
  138.     }
  139.  
  140.     FORR(i, ansR, max(ansL, qryM)) {
  141.         dsu.unite(e[i].u, e[i].v);
  142.         if (dsu.mx_dia > k) {
  143.             ansM = i + 1;
  144.             break;
  145.         }
  146.     }
  147.  
  148.     ans[qryM] = ansM;
  149.     if (ansM == ansR + 1) {
  150.         dsu.roll_back();
  151.         DnC(qryL, qryM - 1, ansL, ansR);
  152.         dsu.snap_shot();
  153.         FOR(i, qryL, qryM) dsu.unite(e[i].u, e[i].v);
  154.         DnC(qryM + 1, qryR, ansL, ansR);
  155.         return;
  156.     }
  157.  
  158.     dsu.roll_back();
  159.  
  160.     dsu.snap_shot();
  161.     FOR(i, ansM + 1, ansR) dsu.unite(e[i].u, e[i].v);
  162.     DnC(qryL, qryM - 1, ansL, ansM);
  163.     dsu.roll_back();
  164.  
  165.     dsu.snap_shot();
  166.     FOR(i, qryL, qryM) dsu.unite(e[i].u, e[i].v);
  167.     DnC(qryM + 1, qryR, ansM, ansR);
  168.     dsu.roll_back();
  169. }
  170.  
  171. void init(void) {
  172.     cin >> n >> k;
  173.     FOR(i, 1, n - 1) {
  174.         int u, v, c;
  175.         cin >> u >> v >> c;
  176.         G[u].pb(mp(v, c));
  177.         G[v].pb(mp(u, c));
  178.         e[i] = {u, v, c};
  179.     }
  180. }
  181.  
  182. void process(void) {
  183.     dfs(1, -1);
  184.     FOR(i, 1, timer) minHigh[i][0] = tour[i];
  185.     FOR(j, 1, 19) FOR(i, 1, timer - MASK(j) + 1) minHigh[i][j] = MIN_HIGH(minHigh[i][j - 1], minHigh[i + MASK(j - 1)][j - 1]);
  186.  
  187.     int r = 0;
  188.     ll res = 0;
  189.     dsu = DSU(n);
  190.     FORR(i, n - 1, 1) {
  191.         dsu.unite(e[i].u, e[i].v);
  192.         if (dsu.mx_dia > k) {
  193.             r = i + 1;
  194.             break;
  195.         }
  196.     }
  197.     if (r == 0) {
  198.         cout << 1LL * n * (n - 1) / 2;
  199.         exit(0);
  200.     } else res += n + 1 - r;
  201.  
  202.     dsu = DSU(n);
  203.     FOR(i, 1, n - 1) ans[i] = n + 1;
  204.     DnC(1, n - 1, 1, n - 1);
  205.     FOR(i, 1, n - 1) res += n + 1 - ans[i];
  206.     cout << res;
  207. }
  208.  
  209. signed main() {
  210.     ios_base::sync_with_stdio(0);
  211.     cin.tie(0); cout.tie(0);
  212.     if (fopen(task".inp", "r")) {
  213.         freopen(task".inp", "r", stdin);
  214.         freopen(task".out", "w", stdout);
  215.     }
  216.     int tc = 1;
  217. //    cin >> tc;
  218.     while(tc--) {
  219.         init();
  220.         process();
  221.     }
  222.     return 0;
  223. }
  224.  
Success #stdin #stdout 0.01s 12396KB
comments (?)
stdin 
Standard input is empty
stdout 
Standard output is empty
https://ideone.com/swQHD3
language:
C++14 (gcc 8.3)
created:
2 days 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!