fork download
copy 
  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
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/d7VGFi
language:
C++14 (gcc 8.3)
created:
7 years 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!