fork download
  1. #include <iostream>
  2. #include <cstdio>
  3. #include <vector>
  4. #include <cstring>
  5.  
  6. using namespace std;
  7.  
  8. #define For(i , a , b) for (int i = a , _b = b ; i <= _b ; ++i)
  9. #define Ford(i , a ,b) for (int i = a , _b = b : i >= _b ; --i)
  10. #define Rep(i , n) for (int i = 0 , _n = n ; i < _n ; ++i)
  11. #define sz(A) ((int)A.size())
  12. #define LL(x) (x << 1)
  13. #define RR(x) ((x << 1) | 1)
  14.  
  15. typedef pair<int , int> pt;
  16.  
  17. const int maxn = 100000 + 123;
  18. int n, color[maxn];
  19. std::vector<int> V[maxn];
  20.  
  21. void ReadData() {
  22. scanf("%d",&n);
  23. For(i, 1, n) scanf("%d", &color[i]);
  24. For(i, 2, n) {
  25. int u, v; scanf("%d%d", &u, &v);
  26. V[u].push_back(v); V[v].push_back(u);
  27. }
  28. }
  29.  
  30. long long res[maxn];
  31. bool deleted[maxn];
  32. int chil[maxn];
  33.  
  34. void dfsSize(const int u, const int prev) {
  35. chil[u] = 1;
  36. for (int v: V[u]) if (v != prev && !deleted[v]) {
  37. dfsSize(v, u);
  38. chil[u] += chil[v];
  39. }
  40. }
  41.  
  42. int newRoot(const int u, const int prev, const int Size) {
  43. for (auto v : V[u]) if (v != prev && !deleted[v]) {
  44. if (chil[v] * 2 >= Size) return newRoot(v, u, Size);
  45. }
  46. return u;
  47. }
  48.  
  49. long long dp[maxn], sum[maxn]; // number of ways for each color
  50. int was[maxn];
  51. int sta[maxn], top = 0;
  52. long long total = 0;
  53. int newChil[maxn];
  54. int imp[maxn];
  55.  
  56. void dfsSizeNew(const int u, const int prev) {
  57. newChil[u] = 1;
  58. for (int v: V[u]) if (v != prev && !deleted[v]) {
  59. dfsSizeNew(v, u);
  60. newChil[u] += newChil[v];
  61. }
  62. }
  63.  
  64. void dfs(const int u, const int prev) {
  65. sta[++top] = u;
  66. if (!was[color[u]]) {
  67. sum[color[u]] += newChil[u];
  68. dp[u] += newChil[u];
  69. imp[u] = newChil[u];
  70. if (prev) total += newChil[u];
  71. } else imp[u] = 0;
  72. was[color[u]]++;
  73. for (auto v : V[u]) if (v != prev && !deleted[v]) {
  74. dfs(v, u);
  75. dp[u] += dp[v];
  76. }
  77. was[color[u]]--;
  78. }
  79.  
  80. int staCol[maxn], top2 = 0, cntCol[maxn];
  81. void preCal(const int u, const int prev) {
  82. if (!cntCol[color[u]]) {
  83. cntCol[color[u]] += imp[u];
  84. staCol[++top2] = color[u];
  85. } else cntCol[color[u]] += imp[u];
  86. for (auto v : V[u]) if (v != prev && !deleted[v]) {
  87. preCal(v, u);
  88. }
  89. }
  90. void resetCol() {
  91. For(i, 1, top2) {
  92. cntCol[staCol[i]] = 0;
  93. }
  94. top2 = 0;
  95. }
  96.  
  97. int totalColor;
  98. void dfsCal(const int u, const int prev, const int anc, const int Size) {
  99. if (!was[color[u]]) {
  100. ++totalColor;
  101. total -= (sum[color[u]] - cntCol[color[u]]);
  102. }
  103. was[color[u]]++;
  104. res[u] += total;
  105. res[u] += 1LL * totalColor * (Size - newChil[anc] );
  106.  
  107. for (auto v : V[u]) if (v != prev && !deleted[v]) {
  108. dfsCal(v, u, anc, Size);
  109. }
  110.  
  111. was[color[u]]--;
  112. if (!was[color[u]]) {
  113. totalColor--;
  114. total += (sum[color[u]] - cntCol[color[u]]);
  115. }
  116. }
  117.  
  118. void Solve(int root, int Size) {
  119. total = 0;
  120. totalColor = 0;
  121.  
  122. dfsSize(root, 0);
  123. int u = newRoot(root, 0, Size);
  124. dfsSizeNew(u, 0);
  125.  
  126. dfs(u, 0);
  127. res[u] += dp[u];
  128. deleted[u] = true;
  129. for (auto v : V[u]) if (!deleted[v]) {
  130. totalColor = 1;
  131. was[color[u]] = 1;
  132. total = dp[u] - dp[v] - newChil[u];
  133. resetCol();
  134. preCal(v, u);
  135. dfsCal(v, u, v, chil[root]);
  136. }
  137.  
  138. For(i, 1, top) {
  139. int u = sta[i];
  140. sum[color[u]] = dp[u] = 0;
  141. was[color[u]] = 0;
  142. imp[u] = 0;
  143. }
  144. top = 0;
  145. for (auto v : V[u]) if (!deleted[v]) {
  146. if (chil[v] < chil[u]) Solve(v, chil[v]); else Solve(v, Size - chil[u]);
  147. }
  148. }
  149.  
  150. void Process() {
  151. Solve(1, n);
  152. For(i, 1, n) cout << res[i] << "\n";
  153. //cout << res[5] << endl;
  154. }
  155.  
  156. int main() {
  157. ios_base::sync_with_stdio(false);
  158. cin.tie(NULL);
  159.  
  160. //freopen("input2.inp" , "r" , stdin);
  161. // freopen("output.out" , "w" , stdout);
  162. ReadData();
  163. Process();
  164.  
  165. return 0;
  166.  
  167. }
Success #stdin #stdout 0s 10216KB
stdin
Standard input is empty
stdout
Standard output is empty