fork download
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. typedef long long ll;
  6. typedef pair<int, int> ii;
  7.  
  8. const int INF = 1e9;
  9. const ll LINF = 1e18;
  10.  
  11. const int N = 2e5 + 5;
  12.  
  13. int n;
  14. int c[N];
  15. vector<int> adj[N];
  16.  
  17. set<int> s[N]; // s[u] = tập hợp các mã màu có trong cây con gốc u
  18. // Lưu ý: sau khi duyệt xong cây có thể s[u] không còn đúng nữa
  19. int ans[N]; // ans[u] = Đáp án cho cây con gốc u
  20.  
  21. void dfs(int u, int p) {
  22. s[u].insert(c[u]);
  23.  
  24. for (int v : adj[u]) {
  25. if (v == p) continue;
  26. dfs(v, u);
  27. // Lúc này đã biết s[v]
  28. // Gộp s[v] với s[u]
  29. if (s[u].size() < s[v].size()) swap(s[u], s[v]); // Chi phí swap 2 CTDL trong C++
  30. // như vector, set, multiset, map,... là O(1)
  31. for (auto it : s[v]) s[u].insert(it);
  32. }
  33.  
  34. // Lúc này đã biết s[u]
  35. ans[u] = s[u].size();
  36. }
  37.  
  38. int main() {
  39. ios::sync_with_stdio(false);
  40. cin.tie(nullptr);
  41. cin >> n;
  42. for (int u = 1; u <= n; u++) cin >> c[u];
  43.  
  44. for (int i = 0; i < n - 1; i++) {
  45. int u, v;
  46. cin >> u >> v;
  47. adj[u].push_back(v);
  48. adj[v].push_back(u);
  49. }
  50.  
  51. dfs(1, -1);
  52.  
  53. for (int u = 1; u <= n; u++) cout << ans[u] << ' ';
  54. }
  55.  
Success #stdin #stdout 0.01s 18188KB
stdin
5
2 3 2 2 1
1 2
1 3
3 4
3 5
stdout
3 1 2 1 1