fork download 
  1. //#pragma GCC optimize("Ofast,unroll-loops")
  2. //#pragma GCC target("avx2,tune=native")
  3. #include <bits/stdc++.h>
  4. using namespace std;
  5.  
  6. #define file "o"
  7. #define ff(i, a, b) for(auto i=(a); i<=(b); ++i)
  8. #define ffr(i, b, a) for(auto i=(b); i>=(a); --i)
  9. #define nl "\n"
  10. #define ss " "
  11. #define pb emplace_back
  12. #define fi first
  13. #define se second
  14. #define sz(s) (int)s.size()
  15. #define all(s) (s).begin(), (s).end()
  16. #define ms(a,x) memset(a, x, sizeof (a))
  17. #define cn continue
  18. #define re exit(0)
  19.  
  20. typedef long long ll;
  21. typedef vector<int> vi;
  22.  
  23. const int maxn = 400000 + 15;
  24. const ll INF = (ll)4e18;
  25.  
  26. int n, m;
  27. int st[maxn], ed[maxn], timer=0;
  28. int st2[maxn], tour2[20][maxn*2], timer2=0;
  29. int h[maxn];
  30. vi g[maxn], ng[maxn], c[maxn], cur;
  31. ll ans[maxn], dff[maxn];
  32. int node_at_tin[maxn];
  33.  
  34. // compute subtree size quickly as ed - st + 1
  35.  
  36. void dfs(int u, int par)
  37. {
  38.     st[u]=++timer;
  39.     st2[u]=++timer2; tour2[0][timer2]=u;
  40.     for(auto v:g[u]) if(v!=par)
  41.     {
  42.         h[v]=h[u]+1;
  43.         dfs(v, u);
  44.         tour2[0][++timer2]=u;
  45.     }
  46.     ed[u]=timer;
  47. }
  48.  
  49. void pre()
  50. {
  51.     int LOG = 19;
  52.     ff(j, 1, LOG) ff(i, 1, timer2-(1<<j)+1)
  53.     {
  54.         int u=tour2[j-1][i], v=tour2[j-1][i+(1<<(j-1))];
  55.         if(h[u]<h[v]) tour2[j][i]=u;
  56.         else tour2[j][i]=v;
  57.     }
  58. }
  59.  
  60. inline int lca(int u, int v)
  61. {
  62.     int l=st2[u], r=st2[v];
  63.     if (l>r) swap(l, r);
  64.     int k=31-__builtin_clz(r-l+1);
  65.     int a = tour2[k][l], b = tour2[k][r-(1<<k)+1];
  66.     if(h[a] < h[b]) return a;
  67.     else return b;
  68. }
  69.  
  70. inline bool cmp_tin(int u, int v)
  71. {
  72.     return st[u] < st[v];
  73. }
  74.  
  75. inline bool child_of(int u, int v)
  76. {
  77.     return st[u] <= st[v] && ed[v] <= ed[u];
  78. }
  79.  
  80. inline void upd(int l, int r, ll val)
  81. {
  82.     if (l>r) return;
  83.     dff[l] += val;
  84.     dff[r+1] -= val;
  85. }
  86.  
  87. void build_virtual()
  88. {
  89.     sort(all(cur), cmp_tin);
  90.     int k = sz(cur);
  91.     // add LCAs of consecutive
  92.     for (int i = 0; i+1 < k; ++i) cur.pb(lca(cur[i], cur[i+1]));
  93.     sort(all(cur), cmp_tin);
  94.     cur.erase(unique(all(cur)), cur.end());
  95.     // clear adjacency
  96.     for (int x : cur) ng[x].clear();
  97.     // build with stack
  98.     vector<int> stck;
  99.     stck.reserve(sz(cur));
  100.     stck.push_back(cur[0]);
  101.     for (int i = 1; i < sz(cur); ++i)
  102.     {
  103.         int v = cur[i];
  104.         while (!child_of(stck.back(), v)) {
  105.             int u = stck.back(); stck.pop_back();
  106.             ng[stck.back()].pb(u);
  107.         }
  108.         stck.push_back(v);
  109.     }
  110.     while (sz(stck) > 1) {
  111.         int u = stck.back(); stck.pop_back();
  112.         ng[stck.back()].pb(u);
  113.     }
  114. }
  115.  
  116. int cnt_in_subtree(const vi &vec, int v)
  117. {
  118.     // vec sorted by tin
  119.     auto lo = lower_bound(vec.begin(), vec.end(), st[v], [&](int node, int value){ return st[node] < value; });
  120.     auto hi = upper_bound(vec.begin(), vec.end(), ed[v], [&](int value, int node){ return value < st[node]; });
  121.     return int(hi - lo);
  122. }
  123.  
  124. int main()
  125. {
  126.     ios::sync_with_stdio(false);
  127.     cin.tie(nullptr); cout.tie(nullptr);
  128.     if (fopen(file".inp","r")) {
  129.         freopen(file".inp","r",stdin);
  130.         freopen(file".out","w",stdout);
  131.     }
  132.     int sub; if(!(cin>>sub)) return 0;
  133.     cin >> n >> m;
  134.     ff(i,1,n){
  135.         int x; cin >> x; c[x].pb(i);
  136.     }
  137.     ff(i,1,n-1){
  138.         int u,v,w; cin >> u >> v >> w;
  139.         g[u].pb(v); g[v].pb(u);
  140.     }
  141.  
  142.     // preprocess
  143.     timer = 0; timer2 = 0;
  144.     h[1] = 0;
  145.     dfs(1, 0);
  146.     pre();
  147.     // build inverse tin
  148.     ff(i,1,n) node_at_tin[st[i]] = i;
  149.  
  150.     // mark array for colored nodes per color
  151.     static int is_col[maxn];
  152.     ms(is_col, 0);
  153.     ms(dff, 0);
  154.     // process each color
  155.     for (int col = 1; col <= n; ++col)
  156.     {
  157.         if (c[col].empty()) continue;
  158.         // vec: original colored nodes sorted by tin
  159.         vi vec = c[col];
  160.         sort(all(vec), cmp_tin);
  161.         // prepare cur = vec (will be appended LCAs inside build_virtual)
  162.         cur.clear();
  163.         for (int x : vec) cur.pb(x);
  164.         // mark colored nodes
  165.         for (int x : vec) is_col[x] = 1;
  166.         // build virtual tree (ng adjacency filled, cur updated)
  167.         build_virtual();
  168.         // iterate virtual nodes in cur and compute contributions
  169.         for (int v : cur)
  170.         {
  171.             // sum sizes of virtual children
  172.             ll sum_ch = 0;
  173.             for (int ch : ng[v]) {
  174.                 int subch = ed[ch] - st[ch] + 1;
  175.                 sum_ch += subch;
  176.             }
  177.             bool v_is_col = is_col[v];
  178.             int sub_v = ed[v] - st[v] + 1;
  179.             if (!v_is_col)
  180.             {
  181.                 // one block: subtree[v] minus children's subtrees
  182.                 ll block_sz = (ll)sub_v - sum_ch;
  183.                 if (block_sz > 0)
  184.                 {
  185.                     ll val = (ll)n - block_sz;
  186.                     // add val to subtree[v]
  187.                     upd(st[v], ed[v], val);
  188.                     // subtract val from each child subtree
  189.                     for (int ch : ng[v]) upd(st[ch], ed[ch], -val);
  190.                 }
  191.             }
  192.             else
  193.             {
  194.                 // v is colored: each neighbor-side that contains 0 colored nodes is a block
  195.                 // iterate original neighbors
  196.                 for (int to : g[v])
  197.                 {
  198.                     if (child_of(to, v)) {
  199.                         // to is ancestor of v -> parent side; shouldn't happen because child_of(to,v) means to is ancestor of v
  200.                         // but in tree adjacency, either to is child of v or parent of v
  201.                     }
  202.                     if (child_of(v, to)) {
  203.                         // 'to' is child of v in original rooted tree
  204.                         int cntcol = cnt_in_subtree(vec, to);
  205.                         int side_sz = ed[to] - st[to] + 1;
  206.                         if (cntcol == 0) {
  207.                             ll val = (ll)n - side_sz;
  208.                             upd(st[to], ed[to], val);
  209.                         }
  210.                     } else {
  211.                         // 'to' is parent side (outside subtree v)
  212.                         int total_col = sz(vec);
  213.                         int inside_v = cnt_in_subtree(vec, v);
  214.                         int outside_col = total_col - inside_v;
  215.                         int side_sz = n - sub_v;
  216.                         if (outside_col == 0 && side_sz > 0) {
  217.                             ll val = (ll)n - side_sz;
  218.                             // outside subtree: two intervals [1, st[v]-1] and [ed[v]+1, n]
  219.                             if (st[v] > 1) upd(1, st[v]-1, val);
  220.                             if (ed[v] < n) upd(ed[v]+1, n, val);
  221.                         }
  222.                     }
  223.                 }
  224.                 // colored node itself: every path to any node contains its own color
  225.                 upd(st[v], st[v], (ll)n);
  226.             }
  227.         }
  228.         // cleanup marks
  229.         for (int x : vec) is_col[x] = 0;
  230.     }
  231.  
  232.     // accumulate diff to ans
  233.     ll curv = 0;
  234.     ff(t,1,n)
  235.     {
  236.         curv += dff[t];
  237.         ans[node_at_tin[t]] = curv;
  238.     }
  239.  
  240.     // print answers
  241.     ff(i,1,n)
  242.     {
  243.         if (i>1) cout << ' ';
  244.         cout << ans[i];
  245.     }
  246.     cout << '\n';
  247.     return 0;
  248. }
  249.  
Success #stdin #stdout 0.02s 51360KB
comments (?)
stdin 
1
5 4
1 2 3 2 3
1 2 9
2 3 7
2 4 2
1 5 2
stdout 
10 9 11 9 12
https://ideone.com/m6oxs1
language:
C++14 (gcc 8.3)
created:
1 day 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!