fork(1) download
  1. #include<bits/stdc++.h>
  2. using namespace std;
  3.  
  4. // Khai báo kích thước mảng tối đa theo đề bài (N <= 2e5)
  5. const int mx = 2e5 + 5;
  6.  
  7. int n;
  8. int c[mx]; // Mảng lưu màu của từng đỉnh
  9. int ans[mx]; // Mảng lưu kết quả (số màu khác nhau trong cây con)
  10. vector<int> adj[mx]; // Danh sách kề để lưu cấu trúc cây
  11. set<int> st[mx]; // Mỗi đỉnh có 1 cái set riêng để lưu các màu trong cây con của nó
  12.  
  13. // Hàm DFS duyệt từ dưới lên (Bottom-up)
  14. void dfs(int u, int p) {
  15. // 1. Khởi tạo: Mỗi đỉnh ban đầu chứa màu của chính nó
  16. st[u].insert(c[u]);
  17.  
  18. // 2. Duyệt qua tất cả các con của u
  19. for (int v : adj[u]) {
  20. if (v != p) { // Tránh quay ngược lại cha
  21. dfs(v, u); // Đệ quy xuống xử lý con trước
  22.  
  23. // 3. Kỹ thuật Small-to-Large (Gộp nhỏ vào lớn)
  24. // So sánh kích thước set của cha (u) và con (v)
  25. // Nếu set của con to hơn set của cha -> Đổi chỗ ngay lập tức
  26. if (st[u].size() < st[v].size()) {
  27. swap(st[u], st[v]); // O(1): Chỉ tráo đổi con trỏ, cực nhanh
  28. }
  29.  
  30. // 4. Thực hiện gộp
  31. // Lúc này đảm bảo st[v] luôn là set nhỏ hơn (hoặc bằng)
  32. // Ta chỉ duyệt qua các phần tử của set nhỏ để bỏ vào set to
  33. for (int color : st[v]) {
  34. st[u].insert(color);
  35. }
  36.  
  37. // Dọn dẹp set con để giải phóng bộ nhớ (không bắt buộc nhưng tốt)
  38. st[v].clear();
  39. }
  40. }
  41.  
  42. // 5. Lưu kết quả
  43. // Sau khi gộp hết tất cả con vào, kích thước set của u chính là số màu khác nhau
  44. ans[u] = st[u].size();
  45. }
  46.  
  47. int main() {
  48. // Tối ưu tốc độ nhập xuất
  49. cin.tie(0)->sync_with_stdio(0);
  50.  
  51. // Mở file nếu chạy trên máy (tiện cho debug)
  52. if (fopen("Distinct_Colors.inp", "r")) {
  53. freopen("Distinct_Colors.inp", "r", stdin);
  54. freopen("Distinct_Colors.out", "w", stdout);
  55. }
  56.  
  57. // Nhập số lượng đỉnh
  58. cin >> n;
  59.  
  60. // Nhập màu cho từng đỉnh
  61. for (int i = 1; i <= n; ++i) cin >> c[i];
  62.  
  63. // Nhập các cạnh và dựng cây
  64. for (int i = 1; i < n; ++i) {
  65. int u, v;
  66. cin >> u >> v;
  67. adj[u].push_back(v);
  68. adj[v].push_back(u);
  69. }
  70.  
  71. // Bắt đầu DFS từ gốc (đỉnh 1), cha của gốc là 0
  72. dfs(1, 0);
  73.  
  74. // In kết quả cho từng đỉnh từ 1 đến n
  75. for (int i = 1; i <= n; ++i) {
  76. cout << ans[i] << " ";
  77. }
  78.  
  79. return 0;
  80. }
  81.  
Success #stdin #stdout 0.01s 18940KB
stdin
Standard input is empty
stdout
Standard output is empty