fork download
copy 
  1. #include <iostream>
  2. #include <vector>
  3.  
  4. using namespace std;
  5.  
  6. const int MAXN = 100005;
  7.  
  8. int n;
  9. int color[MAXN];
  10. vector<int> adj[MAXN];
  11. bool is_centroid[MAXN];
  12. int sz[MAXN];
  13.  
  14. int freq[MAXN], freq_b[MAXN], cnt[MAXN];
  15. long long SumFreq, SumFreq_b;
  16. long long ans[MAXN];
  17.  
  18. // Tính kích thước các cây con
  19. void get_sizes(int u, int p) {
  20.     sz[u] = 1;
  21.     for (int v : adj[u]) {
  22.         if (v != p && !is_centroid[v]) {
  23.             get_sizes(v, u);
  24.             sz[u] += sz[v];
  25.         }
  26.     }
  27. }
  28.  
  29. // Tìm đỉnh trọng tâm (Centroid)
  30. int get_centroid(int u, int p, int total) {
  31.     for (int v : adj[u]) {
  32.         if (v != p && !is_centroid[v] && sz[v] > total / 2) {
  33.             return get_centroid(v, u, total);
  34.         }
  35.     }
  36.     return u;
  37. }
  38.  
  39. // Tính kích thước tương đối so với C và mảng freq toàn cục cho Centroid C
  40. void get_sz_and_freq(int u, int p) {
  41.     sz[u] = 1;
  42.     cnt[color[u]]++;
  43.     bool first = (cnt[color[u]] == 1);
  44.     for (int v : adj[u]) {
  45.         if (v != p && !is_centroid[v]) {
  46.             get_sz_and_freq(v, u);
  47.             sz[u] += sz[v];
  48.         }
  49.     }
  50.     if (first) {
  51.         freq[color[u]] += sz[u];
  52.         SumFreq += sz[u];
  53.     }
  54.     cnt[color[u]]--;
  55. }
  56.  
  57. // Tính mảng freq_b riêng cho nhánh con
  58. void get_freq_b(int u, int p) {
  59.     cnt[color[u]]++;
  60.     bool first = (cnt[color[u]] == 1);
  61.     for (int v : adj[u]) {
  62.         if (v != p && !is_centroid[v]) {
  63.             get_freq_b(v, u);
  64.         }
  65.     }
  66.     if (first) {
  67.         freq_b[color[u]] += sz[u];
  68.         SumFreq_b += sz[u];
  69.     }
  70.     cnt[color[u]]--;
  71. }
  72.  
  73. // Duyệt cây trong nhánh để cộng dồn kết quả vào ans[u]
  74. void solve_branch(int u, int p, long long sum_freq, long long sum_freq_b, int path_len, int tot_nodes, int sz_root_b) {
  75.     cnt[color[u]]++;
  76.     bool first = (cnt[color[u]] == 1);
  77.     if (first) {
  78.         sum_freq += freq[color[u]];
  79.         sum_freq_b += freq_b[color[u]];
  80.         path_len++;
  81.     }
  82.  
  83.     // Áp dụng công thức đóng góp của các đỉnh nằm ngoài nhánh hiện tại
  84.     ans[u] += 1LL * path_len * (tot_nodes - sz_root_b) + SumFreq - SumFreq_b - sum_freq + sum_freq_b;
  85.  
  86.     for (int v : adj[u]) {
  87.         if (v != p && !is_centroid[v]) {
  88.             solve_branch(v, u, sum_freq, sum_freq_b, path_len, tot_nodes, sz_root_b);
  89.         }
  90.     }
  91.  
  92.     cnt[color[u]]--;
  93. }
  94.  
  95. // Xóa trạng thái của mảng tần số (tránh dùng memset toàn bộ để đảm bảo O(N log N))
  96. void clear_freq_b(int u, int p) {
  97.     freq_b[color[u]] = 0;
  98.     for (int v : adj[u]) {
  99.         if (v != p && !is_centroid[v]) {
  100.             clear_freq_b(v, u);
  101.         }
  102.     }
  103. }
  104.  
  105. void clear_freq(int u, int p) {
  106.     freq[color[u]] = 0;
  107.     for (int v : adj[u]) {
  108.         if (v != p && !is_centroid[v]) {
  109.             clear_freq(v, u);
  110.         }
  111.     }
  112. }
  113.  
  114. // Quá trình Centroid Decomposition
  115. void decompose(int u) {
  116.     get_sizes(u, 0);
  117.     int total = sz[u];
  118.     int C = get_centroid(u, 0, total);
  119.  
  120.     // Bước 1: Khởi tạo dữ liệu liên quan đến C
  121.     SumFreq = 0;
  122.     get_sz_and_freq(C, 0);
  123.     ans[C] += SumFreq;
  124.  
  125.     // Bước 2: Duyệt từng nhánh con nối với C
  126.     for (int root_b : adj[C]) {
  127.         if (!is_centroid[root_b]) {
  128.             SumFreq_b = 0;
  129.  
  130.             // Màu của đỉnh Centroid C luôn thuộc mọi đường đi trong thành phần liên thông
  131.             cnt[color[C]]++;
  132.             freq_b[color[C]] += sz[root_b];
  133.             SumFreq_b += sz[root_b];
  134.  
  135.             get_freq_b(root_b, C);
  136.  
  137.             long long init_sum_freq = freq[color[C]];
  138.             long long init_sum_freq_b = freq_b[color[C]];
  139.  
  140.             // Tính toán và update mảng ans cho tất cả các đỉnh nhánh hiện tại
  141.             solve_branch(root_b, C, init_sum_freq, init_sum_freq_b, 1, sz[C], sz[root_b]);
  142.  
  143.             cnt[color[C]]--;
  144.  
  145.             // Xóa bộ đệm tần số riêng biệt của nhánh đó
  146.             clear_freq_b(root_b, C);
  147.             freq_b[color[C]] = 0;
  148.         }
  149.     }
  150.  
  151.     clear_freq(C, 0);
  152.     is_centroid[C] = true;
  153.  
  154.     // Phân rã đệ quy trên các thành phần liên thông còn lại
  155.     for (int v : adj[C]) {
  156.         if (!is_centroid[v]) {
  157.             decompose(v);
  158.         }
  159.     }
  160. }
  161.  
  162. int main() {
  163.     // Tối ưu I/O cho C++
  164.     ios_base::sync_with_stdio(false);
  165.     cin.tie(NULL);
  166.  
  167.     if (!(cin >> n)) return 0;
  168.  
  169.     for (int i = 1; i <= n; i++) {
  170.         cin >> color[i];
  171.     }
  172.  
  173.     for (int i = 1; i < n; i++) {
  174.         int u, v;
  175.         cin >> u >> v;
  176.         adj[u].push_back(v);
  177.         adj[v].push_back(u);
  178.     }
  179.  
  180.     decompose(1);
  181.  
  182.     for (int i = 1; i <= n; i++) {
  183.         cout << ans[i] << (i == n ? "" : " ");
  184.     }
  185.     cout << "\n";
  186.  
  187.     return 0;
  188. }
  189.  
Success #stdin #stdout 0.01s 7016KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/PfnvGh
language:
C++14 (gcc 8.3)
created:
4 hours 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!