fork download
copy 
  1. #pragma GCC optimize("O3,unroll-loops")
  2. #include <bits/stdc++.h>
  3. #define ll long long
  4. #define itachi ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
  5. #define fi first
  6. #define se second
  7. #define maxn 500005
  8.  
  9. using namespace std;
  10.  
  11. int32_t n;
  12. int32_t sz[maxn];
  13. bool banned[maxn];
  14. int32_t cnt[maxn];
  15. int32_t delta_only[maxn], delta_all[maxn];
  16. int32_t delta_only_inside[maxn], delta_all_inside[maxn];
  17. vector<pair<int32_t, int32_t>> adj[maxn];
  18.  
  19. long long ans = 0;
  20. int32_t sum_only = 0;
  21. int32_t sum_only_inside = 0;
  22. vector<int32_t> usedDelta, usedInside;
  23. int32_t current_root = 0;
  24.  
  25. void get_sz(int32_t u, int32_t p) {
  26.     sz[u] = 1;
  27.     for (const auto& [v, w] : adj[u]) {
  28.         if (v != p && !banned[v]) {
  29.             get_sz(v, u);
  30.             sz[u] += sz[v];
  31.         }
  32.     }
  33. }
  34.  
  35. void countChild(int32_t u, int32_t p, int32_t c) {
  36.     sz[u] = 1;
  37.     bool first = false;
  38.  
  39.     if (u != current_root) first = (++cnt[c] == 1);
  40.  
  41.     for (const auto& [v, w] : adj[u]) {
  42.         if (v == p || banned[v]) continue;
  43.         countChild(v, u, w);
  44.         sz[u] += sz[v];
  45.     }
  46.  
  47.     if (u != current_root) {
  48.         if (first) {
  49.             sum_only += sz[u];
  50.             if (delta_all[c] == 0) usedDelta.push_back(c);
  51.             delta_all[c] += sz[u];
  52.             delta_only[c] += sz[u];
  53.         }
  54.         else if (cnt[c] == 2) {
  55.             sum_only -= sz[u];
  56.             delta_only[c] -= sz[u];
  57.         }
  58.         --cnt[c];
  59.     }
  60. }
  61.  
  62. int32_t find_cen(int32_t u, int32_t p, int32_t S) {
  63.     for (const auto& [v, w] : adj[u]) {
  64.         if (v != p && !banned[v] && sz[v] > S / 2) {
  65.             return find_cen(v, u, S);
  66.         }
  67.     }
  68.     return u;
  69. }
  70.  
  71. void cntInside(int32_t u, int32_t p, int32_t c) {
  72.     bool first = (++cnt[c] == 1);
  73.  
  74.     if (first) {
  75.         sum_only_inside += sz[u];
  76.         if (delta_all_inside[c] == 0) usedInside.push_back(c);
  77.         delta_only_inside[c] += sz[u];
  78.         delta_all_inside[c] += sz[u];
  79.     }
  80.     else if (cnt[c] == 2) {
  81.         sum_only_inside -= sz[u];
  82.         delta_only_inside[c] -= sz[u];
  83.     }
  84.  
  85.     for (const auto& [v, w] : adj[u]) {
  86.         if (v == p || banned[v]) continue;
  87.         cntInside(v, u, w);
  88.     }
  89.  
  90.     --cnt[c];
  91. }
  92.  
  93. void addRootPaths(int32_t u, int32_t p, int32_t distinct, int32_t c) {
  94.     bool first = (++cnt[c] == 1);
  95.     if (first) distinct++;
  96.     else if (cnt[c] == 2) distinct--;
  97.  
  98.     ans += 2LL * distinct;
  99.  
  100.     for (const auto& [v, w] : adj[u]) {
  101.         if (v == p || banned[v]) continue;
  102.         addRootPaths(v, u, distinct, w);
  103.     }
  104.     --cnt[c];
  105. }
  106.  
  107. void dfs(int32_t u, int32_t p, int32_t C, int32_t only_cur, int32_t all_cur, int32_t only_out, int32_t distinctlen, int32_t ver_out) {
  108.     bool first = (++cnt[C] == 1);
  109.  
  110.     if (first) {
  111.         distinctlen++;
  112.         all_cur += delta_all[C] - delta_all_inside[C];
  113.         only_cur += delta_only[C] - delta_only_inside[C];
  114.     }
  115.     else if (cnt[C] == 2) {
  116.         distinctlen--;
  117.         all_cur -= delta_all[C] - delta_all_inside[C];
  118.     }
  119.     ans += 1LL * distinctlen * ver_out - all_cur - only_cur + only_out;
  120.  
  121.     for (const auto& [v, w] : adj[u]) {
  122.         if (v == p || banned[v]) continue;
  123.         dfs(v, u, w, only_cur, all_cur, only_out, distinctlen, ver_out);
  124.     }
  125.  
  126.     --cnt[C];
  127. }
  128.  
  129. void solve(int32_t u) {
  130.     get_sz(u, 0);
  131.     int32_t root = find_cen(u, 0, sz[u]);
  132.  
  133.     current_root = root;
  134.     sum_only = 0;
  135.     countChild(root, 0, 0);
  136.     int32_t all_node = sz[root];
  137.  
  138.     for (const auto& [v, w] : adj[root]) {
  139.         if (banned[v]) continue;
  140.         addRootPaths(v, root, 0, w);
  141.     }
  142.  
  143.     for (const auto& [v, w] : adj[root]) {
  144.         if (banned[v]) continue;
  145.  
  146.         sum_only_inside = 0;
  147.         cntInside(v, root, w);
  148.         int32_t ver_out = all_node - sz[v] - 1;
  149.         int32_t only_out = sum_only - sum_only_inside;
  150.  
  151.         dfs(v, root, w, 0, 0, only_out, 0, ver_out);
  152.  
  153.         for (int32_t c : usedInside) {
  154.             delta_all_inside[c] = 0;
  155.             delta_only_inside[c] = 0;
  156.         }
  157.         usedInside.clear();
  158.     }
  159.  
  160.     banned[root] = 1;
  161.     for (int32_t c : usedDelta) {
  162.         delta_all[c] = 0;
  163.         delta_only[c] = 0;
  164.     }
  165.     usedDelta.clear();
  166.     for (const auto& [v, w] : adj[root]) {
  167.         if (!banned[v]) solve(v);
  168.     }
  169. }
  170.  
  171. int main() {
  172.     itachi
  173.     cin>>n;
  174.     for (int32_t i = 1; i < n; i++) {
  175.         int32_t u, v, w;
  176.         cin >> u >> v >> w;
  177.         adj[u].push_back({v, w});
  178.         adj[v].push_back({u, w});
  179.     }
  180.  
  181.     solve(1);
  182.     cout << ans / 2;
  183.     return 0;
  184. }
  185.  
Success #stdin #stdout 0.01s 18548KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/LyYyEo
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!