fork download
copy 
  1. // ~~ icebear ~~
  2. #include <bits/stdc++.h>
  3. using namespace std;
  4. #define int long long
  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 "color"
  33. /*END OF TEMPLATE. ICEBEAR AND THE CAT WILL WIN VOI26 */
  34.  
  35. const int MOD = 1e9 + 7;
  36. const int inf = 1e9 + 27092008;
  37. const ll INF = 1e18 + 27092008;
  38. const int N = 2e5 + 5;
  39. int numNode, numEdge;
  40. vector<int> G[N];
  41. int color[N], cnt[N], diff_color;
  42. vector<int> colors;
  43.  
  44. void add(int u) {
  45.     if (++cnt[color[u]] == 1) {
  46.         diff_color++;
  47.         colors.pb(color[u]);
  48.     }
  49. }
  50.  
  51. void del(int u) {
  52.     if (--cnt[color[u]] == 0) {
  53.         diff_color--;
  54.         colors.pop_back();
  55.     }
  56. }
  57.  
  58. namespace Subtask1 {
  59.     bool check() {
  60.         return numNode <= 5000;
  61.     }
  62.     int ans[N];
  63.     void dfs(int root, int u, int par) {
  64.         add(u);
  65.         ans[root] += diff_color;
  66.         for(int v : G[u]) if (v != par) dfs(root, v, u);
  67.         del(u);
  68.     }
  69.  
  70.     void solve() {
  71.         FOR(i, 1, numNode) {
  72.             dfs(i, i, -1);
  73.             cout << ans[i] << ' ';
  74.         }
  75.     }
  76. }
  77.  
  78. namespace Subtask5 {
  79.     int sz[N];
  80.     bool isCentroid[N];
  81.     int ans[N];
  82.  
  83.     int get_sz(int u, int par) {
  84.         sz[u] = 1;
  85.         for(int v : G[u]) if (v != par && !isCentroid[v])
  86.             sz[u] += get_sz(v, u);
  87.         return sz[u];
  88.     }
  89.  
  90.     int findCentroid(int u, int par, const int &subtree) {
  91.         for(int v : G[u]) if (v != par && !isCentroid[v] && sz[v] * 2 > subtree)
  92.             return findCentroid(v, u, subtree);
  93.         return u;
  94.     }
  95.  
  96.     int countNode, sumDiff, contain[N], f[N];
  97.     void prepare(int u, int par) {
  98.         add(u);
  99.         f[u] = diff_color;
  100.         sumDiff += diff_color;
  101.         countNode++;
  102.         for(int &c : colors) contain[c]++;
  103.         for(int v : G[u]) if (v != par && !isCentroid[v]) prepare(v, u);
  104.         del(u);
  105.     }
  106.  
  107.     void updateTree(int u, int par, int sign) {
  108.         add(u);
  109.         countNode += sign;
  110.         sumDiff += diff_color * sign;
  111.         for(int &c : colors) contain[c] += sign;
  112.         for(int v : G[u]) if (v != par && !isCentroid[v])
  113.             updateTree(v, u, sign);
  114.         del(u);
  115.     }
  116.  
  117.     void solve(int u, int par) {
  118.         add(u);
  119.         ans[u] += f[u] * countNode + sumDiff;
  120.         for(int &c : colors) ans[u] -= contain[c]; // outside subtree
  121.         for(int v : G[u]) if (v != par && !isCentroid[v]) solve(v, u);
  122.         del(u);
  123.     }
  124.  
  125.     void reset(int u, int par) {
  126.         contain[color[u]] = 0;
  127.         for(int v : G[u]) if (v != par && !isCentroid[v])
  128.             reset(v, u);
  129.     }
  130.  
  131.     void buildCentroid(int u) {
  132.         int centroid = findCentroid(u, -1, get_sz(u, -1));
  133.         isCentroid[centroid] = true;
  134.  
  135.         sumDiff = countNode = 0;
  136.  
  137.         add(centroid);
  138.         prepare(centroid, -1);
  139.  
  140.         for(int &v : G[centroid]) if (!isCentroid[v]) {
  141.             updateTree(v, centroid, -1);
  142.             solve(v, centroid);
  143.             updateTree(v, centroid, +1);
  144.         }
  145.  
  146.         ans[centroid] += sumDiff;
  147.         del(centroid);
  148.         reset(centroid, -1);
  149.  
  150.         for(int &v : G[centroid]) if (!isCentroid[v])
  151.             buildCentroid(v);
  152.     }
  153.  
  154.     void solve() {
  155.         buildCentroid(1);
  156.         FOR(i, 1, numNode) cout << ans[i] << ' ';
  157.     }
  158. }
  159.  
  160. void init(void) {
  161.     int subtask; cin >> subtask;
  162.     cin >> numNode >> numEdge;
  163.     FOR(i, 1, numNode) cin >> color[i];
  164.     FOR(i, 1, numEdge) {
  165.         int u, v, w;
  166.         cin >> u >> v >> w;
  167.         G[u].pb(v);
  168.         G[v].pb(u);
  169.     }
  170. }
  171.  
  172. void process(void) {
  173.     if (Subtask1 :: check()) Subtask1 :: solve();
  174.     else Subtask5 :: solve();
  175. }
  176.  
  177. signed main() {
  178.     ios_base::sync_with_stdio(0);
  179.     cin.tie(0); cout.tie(0);
  180.     if (fopen(task".inp", "r")) {
  181.         freopen(task".inp", "r", stdin);
  182.         freopen(task".out", "w", stdout);
  183.     }
  184.     int tc = 1;
  185. //    cin >> tc;
  186.     while(tc--) {
  187.         init();
  188.         process();
  189.     }
  190.     return 0;
  191. }
  192.  
Success #stdin #stdout 0.01s 11452KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/BLvQnk
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!