fork download 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. using ll = long long;
  4.  
  5. int main(){
  6.     ios::sync_with_stdio(false);
  7.     cin.tie(nullptr);
  8.     int theta;
  9.     if (!(cin >> theta)) return 0;
  10.     int n,m;
  11.     cin >> n >> m;
  12.     vector<int> color(n+1);
  13.     for (int i = 1; i <= n; ++i) cin >> color[i];
  14.     vector<vector<int>> g(n+1);
  15.     for (int i = 0; i < m; ++i){
  16.         int u,v,w; cin >> u >> v >> w;
  17.         g[u].push_back(v);
  18.         g[v].push_back(u);
  19.     }
  20.  
  21.     // Preprocess: root at 1, compute tin/tout, up table, subtree sizes
  22.     int LOG = 1;
  23.     while ((1<<LOG) <= n) ++LOG;
  24.     vector<vector<int>> up(LOG, vector<int>(n+1));
  25.     vector<int> tin(n+1), tout(n+1), depth(n+1), sub(n+1);
  26.     int timer = 0;
  27.     // iterative DFS to compute tin/tout, parent, depth, subtree
  28.     vector<tuple<int,int,int>> st;
  29.     st.reserve(2*n);
  30.     st.emplace_back(1, 0, 0);
  31.     while (!st.empty()){
  32.         auto [v,p,vis] = st.back(); st.pop_back();
  33.         if (!vis){
  34.             tin[v] = ++timer;
  35.             up[0][v] = (p==0? v : p);
  36.             depth[v] = (p==0? 0 : depth[p]+1);
  37.             st.emplace_back(v,p,1);
  38.             for (int to : g[v]) if (to != p) st.emplace_back(to, v, 0);
  39.         } else {
  40.             int s = 1;
  41.             for (int to : g[v]) if (to != up[0][v]) s += sub[to];
  42.             sub[v] = s;
  43.             tout[v] = timer;
  44.         }
  45.     }
  46.     for (int k = 1; k < LOG; ++k){
  47.         for (int v = 1; v <= n; ++v) up[k][v] = up[k-1][ up[k-1][v] ];
  48.     }
  49.     auto is_anc = [&](int a,int b)->bool{ return tin[a] <= tin[b] && tout[b] <= tout[a]; };
  50.     auto lca = [&](int a,int b){
  51.         if (is_anc(a,b)) return a;
  52.         if (is_anc(b,a)) return b;
  53.         for (int k = LOG-1; k >= 0; --k) if (!is_anc(up[k][a], b)) a = up[k][a];
  54.         return up[0][a];
  55.     };
  56.  
  57.     // group nodes by color
  58.     vector<vector<int>> by_color(n+1);
  59.     for (int v = 1; v <= n; ++v) by_color[color[v]].push_back(v);
  60.  
  61.     // difference array on Euler to add to subtree intervals
  62.     vector<ll> diff(n+3, 0);
  63.     auto add_range = [&](int L,int R,ll val){
  64.         if (L > R) return;
  65.         diff[L] += val;
  66.         diff[R+1] -= val;
  67.     };
  68.     vector<int> node_at_tin(n+1);
  69.     for (int v = 1; v <= n; ++v) node_at_tin[tin[v]] = v;
  70.  
  71.     // temporary containers
  72.     vector<int> nodes; nodes.reserve(1024);
  73.     vector<int> stk2; stk2.reserve(1024);
  74.     vector<vector<int>> vt(n+1);
  75.  
  76.     for (int col = 1; col <= n; ++col){
  77.         auto &vec = by_color[col];
  78.         if (vec.empty()) continue;
  79.         // sort colored nodes by tin (for binary searching counts)
  80.         sort(vec.begin(), vec.end(), [&](int a,int b){ return tin[a] < tin[b]; });
  81.  
  82.         // build virtual tree node set: colored nodes + LCAs
  83.         nodes.clear();
  84.         for (int x : vec) nodes.push_back(x);
  85.         int origin_sz = nodes.size();
  86.         for (int i = 0; i + 1 < origin_sz; ++i){
  87.             nodes.push_back(lca(nodes[i], nodes[i+1]));
  88.         }
  89.         sort(nodes.begin(), nodes.end(), [&](int a,int b){ if (tin[a]==tin[b]) return a<b; return tin[a] < tin[b]; });
  90.         nodes.erase(unique(nodes.begin(), nodes.end()), nodes.end());
  91.  
  92.         // build virtual tree adjacency
  93.         for (int x : nodes) vt[x].clear();
  94.         stk2.clear();
  95.         for (int x : nodes){
  96.             while (!stk2.empty() && !is_anc(stk2.back(), x)) stk2.pop_back();
  97.             if (!stk2.empty()) vt[stk2.back()].push_back(x);
  98.             stk2.push_back(x);
  99.         }
  100.  
  101.         // helper: count how many colored nodes lie inside subtree v
  102.         auto cnt_in_subtree = [&](int v)->int{
  103.             auto lo = lower_bound(vec.begin(), vec.end(), tin[v], [&](int node,int value){ return tin[node] < value; });
  104.             auto hi = upper_bound(vec.begin(), vec.end(), tout[v], [&](int value,int node){ return value < tin[node]; });
  105.             return int(hi - lo);
  106.         };
  107.  
  108.         // process virtual nodes
  109.         for (int v : nodes){
  110.             int is_colored = (binary_search(vec.begin(), vec.end(), v, [&](int a,int b){ return tin[a] < tin[b]; })
  111.                               || (tin[vec[ lower_bound(vec.begin(), vec.end(), tin[v], [&](int node,int value){ return tin[node] < value; }) - vec.begin() ]] == tin[v] )
  112.                               ) ? 0 : 0;
  113.             // simpler check:
  114.             bool v_is_col = false;
  115.             auto itv = lower_bound(vec.begin(), vec.end(), tin[v], [&](int node,int value){ return tin[node] < value; });
  116.             if (itv != vec.end() && tin[*itv] == tin[v]) v_is_col = true;
  117.  
  118.             // sum subtree sizes of virtual children
  119.             ll sum_ch_sub = 0;
  120.             for (int c : vt[v]) sum_ch_sub += sub[c];
  121.  
  122.             if (!v_is_col){
  123.                 // remainder in subtree[v] not covered by children in VT -> a single connected block
  124.                 ll block_sz = (ll)sub[v] - sum_ch_sub;
  125.                 if (block_sz > 0){
  126.                     ll val = (ll)n - block_sz;
  127.                     // add val to subtree[v]
  128.                     add_range(tin[v], tout[v], val);
  129.                     // subtract from each vt-child subtree
  130.                     for (int c : vt[v]) add_range(tin[c], tout[c], -val);
  131.                 }
  132.             } else {
  133.                 // v is a colored node: each original neighbor side that contains 0 colored nodes becomes a component
  134.                 // we iterate all neighbors in original tree (g[v])
  135.                 for (int to : g[v]){
  136.                     int side_sz;
  137.                     int cnt_col_side;
  138.                     if (up[0][to] == v){ // to is child of v
  139.                         side_sz = sub[to];
  140.                         cnt_col_side = cnt_in_subtree(to);
  141.                         if (cnt_col_side == 0){
  142.                             ll val = (ll)n - side_sz;
  143.                             add_range(tin[to], tout[to], val);
  144.                         }
  145.                     } else if (up[0][v] == to){ // to is parent of v
  146.                         side_sz = n - sub[v];
  147.                         // number of colored nodes outside subtree[v] = total colored - cnt_in_subtree(v)
  148.                         int total_col = (int)vec.size();
  149.                         int inside = cnt_in_subtree(v);
  150.                         cnt_col_side = total_col - inside;
  151.                         if (cnt_col_side == 0){
  152.                             ll val = (ll)n - side_sz;
  153.                             // nodes outside subtree[v] are two intervals: [1, tin[v]-1] and [tout[v]+1, n]
  154.                             if (tin[v] > 1) add_range(1, tin[v]-1, val);
  155.                             if (tout[v] < n) add_range(tout[v]+1, n, val);
  156.                         }
  157.                     } else {
  158.                         // should not happen in tree adjacency
  159.                     }
  160.                 }
  161.                 // colored node itself contributes n (path from v to any vertex contains its color)
  162.                 add_range(tin[v], tin[v], n);
  163.             }
  164.         }
  165.     }
  166.  
  167.     // accumulate diff to answer
  168.     vector<ll> ans(n+1, 0);
  169.     ll cur = 0;
  170.     for (int t = 1; t <= n; ++t){
  171.         cur += diff[t];
  172.         ans[node_at_tin[t]] = cur;
  173.     }
  174.     for (int i = 1; i <= n; ++i){
  175.         if (i>1) cout << ' ';
  176.         cout << ans[i];
  177.     }
  178.     cout << '\n';
  179.     return 0;
  180. }
  181.  
Success #stdin #stdout 0s 5316KB
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/SXsCgq
language:
C++14 (gcc 8.3)
created:
3 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!