fork download
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. const int maxn = 4e5+5;
  5. int n, m;
  6. vector<int> g[maxn];
  7. int color[maxn];
  8. long long ans[maxn];
  9.  
  10. // DSU on Tree variables
  11. map<int,int>* cnt[maxn]; // pointer to map<int,int> cho subtree
  12. int subtree[maxn]; // kích thước subtree
  13.  
  14. // DFS chuẩn bị kích thước subtree
  15. void dfs_size(int u, int p) {
  16. subtree[u] = 1;
  17. for(auto v : g[u]) if(v != p) {
  18. dfs_size(v,u);
  19. subtree[u] += subtree[v];
  20. }
  21. }
  22.  
  23. // DFS DSU on Tree
  24. void dfs(int u, int p, bool keep) {
  25. int mx = -1, bigChild = -1;
  26. for(auto v : g[u]) if(v != p && subtree[v] > mx) {
  27. mx = subtree[v];
  28. bigChild = v;
  29. }
  30.  
  31. for(auto v : g[u]) if(v != p && v != bigChild)
  32. dfs(v,u,false); // xử lý nhánh nhỏ trước
  33.  
  34. if(bigChild != -1) dfs(bigChild,u,true), cnt[u] = cnt[bigChild];
  35. else cnt[u] = new map<int,int>();
  36.  
  37. (*cnt[u])[color[u]]++;
  38.  
  39. for(auto v : g[u]) if(v != p && v != bigChild) {
  40. for(auto &pr : *cnt[v]) (*cnt[u])[pr.first] += pr.second;
  41. }
  42.  
  43. // Tính ans[u] = tổng số màu trên mọi đường từ u tới các đỉnh subtree
  44. // Ở đây mỗi màu count 1 cho mỗi đỉnh, ta cần số màu distinct = cnt[u].size()
  45. ans[u] = 0;
  46. for(auto &pr : *cnt[u]) {
  47. ans[u] += pr.second;
  48. }
  49.  
  50. if(!keep) delete cnt[u]; // xóa nhánh nhỏ
  51. }
  52.  
  53. int main() {
  54. ios::sync_with_stdio(false);
  55. cin.tie(nullptr);
  56.  
  57. cin >> n >> m; // m = n-1
  58. for(int i=1;i<=n;i++) cin >> color[i];
  59.  
  60. for(int i=0;i<m;i++){
  61. int u,v,w; cin >> u >> v >> w;
  62. g[u].push_back(v);
  63. g[v].push_back(u);
  64. }
  65.  
  66. dfs_size(1,0);
  67. dfs(1,0,true);
  68.  
  69. for(int i=1;i<=n;i++){
  70. cout << ans[i] << " ";
  71. }
  72. cout << "\n";
  73. return 0;
  74. }
  75.  
Success #stdin #stdout 0.01s 19844KB
stdin
46 45
23 21 20 9 26 10 26 41 15 41 21 41 23 38 43 31 3 14 7 23 43 44 27 14 18 42 21 37 45 25 11 24 42 20 16 27 25 43 46 10 11 38 30 32 41 39
2 1 86
3 1 56
4 3 38
5 2 98
6 1 31
7 2 91
8 3 38
9 4 99
10 8 51
11 3 8
12 10 99
13 1 52
14 6 84
15 7 74
16 7 48
17 11 40
18 10 58
19 10 83
20 17 84
21 5 81
22 2 83
23 18 92
24 6 31
25 4 22
26 16 90
27 26 72
28 20 97
29 16 64
30 22 72
31 17 86
32 4 49
33 26 1
34 3 58
35 4 72
36 20 26
37 13 50
38 11 75
39 18 3
40 10 69
41 40 84
42 39 64
43 9 42
44 13 60
45 17 84
46 13 18
stdout
3720 92 464 24 2 4 22 57 2 56 42 1 8 1 1 10 20 6 1 4 1 2 1 1 1 4 1 1 1 1 1 1 1 1 1 1 1 1 2 2 1 1 1 1 1 1