fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. #define int              long long int
  4. #define double           long double
  5. inline int power(int a, int b) {
  6.     int x = 1;
  7.     while (b) {
  8.         if (b & 1) x *= a;
  9.         a *= a;
  10.         b >>= 1;
  11.     }
  12.     return x;
  13. }
  14.  
  15.  
  16. const int M = 1000000007;
  17. const int N = 3e5+9;
  18. const int INF = 2e9+1;
  19. const int LINF = 2000000000000000001;
  20.  
  21. //_ ***************************** START Below *******************************
  22.  
  23.  
  24.  
  25. vector<vector<int>> edges;
  26. vector<int> coins;
  27.  
  28. void consistency(int n) {
  29.  
  30. 	vector<vector<int>> graph(n);
  31. 	for(int i=0; i<n-1; i++){
  32. 		int x = edges[i][0], y = edges[i][1];
  33. 		graph[x].push_back(y);
  34. 		graph[y].push_back(x);
  35. 	}
  36.  
  37. 	int removed = 0;
  38. 	vector<int> degree(n);
  39. 	queue<int> q;
  40.  
  41.  
  42. 	//* Step 0 : Multi Source BFS to recursively delete all leave nodes 
  43. 	for(int i=0; i<n; i++){
  44. 		degree[i] = graph[i].size();
  45. 		if(degree[i] == 1 && coins[i] == 0){
  46. 			q.push(i);
  47. 		}
  48. 	}
  49.  
  50. 	while(!q.empty()){
  51. 		removed++;
  52. 		auto ch = q.front(); q.pop();
  53.  
  54. 		int parent = -1;
  55.  
  56. 		for(auto& par : graph[ch]){
  57. 			degree[par]--;
  58. 			parent = par;
  59. 			if(degree[par] == 1 && coins[par] == 0){
  60. 				q.push(par);
  61. 			}
  62. 		}
  63.  
  64. 		if(parent != -1){
  65. 			auto ch_to_delete = find(graph[parent].begin(), graph[parent].end(), ch);
  66. 			graph[parent].erase(ch_to_delete);
  67.  
  68. 			auto par_to_delete = find(graph[ch].begin(), graph[ch].end(), parent);
  69. 			graph[ch].erase(par_to_delete );
  70. 		}
  71. 	}
  72.  
  73.  
  74. 	//* Step1 : Multi source BFS to delete all leaf nodes with 1's 
  75. 	//* After deleting all 0's leaf nodes, we are guranteed to have leaf nodes with all 1's
  76. 	for(int i=0; i<n; i++){
  77. 		degree[i] = graph[i].size();
  78. 		if(degree[i] == 1 ){
  79. 			q.push(i);
  80. 		}
  81. 	}
  82.  
  83. 	while(!q.empty()){
  84. 		removed++;
  85. 		auto ch = q.front(); q.pop();
  86.  
  87. 		int parent = -1;
  88.  
  89. 		for(auto& par : graph[ch]){
  90. 			degree[par]--;
  91. 			parent = par;
  92. 		}
  93.  
  94. 		if(parent != -1){
  95. 			auto ch_to_delete = find(graph[parent].begin(), graph[parent].end(), ch);
  96. 			graph[parent].erase(ch_to_delete);
  97.  
  98. 			auto par_to_delete = find(graph[ch].begin(), graph[ch].end(), parent);
  99. 			graph[ch].erase(par_to_delete );
  100. 		}
  101. 	}
  102.  
  103.  
  104. 	//* Step 2 : Multi source BFS to delete all leaf nodes with 1's 
  105. 	//* After deleting all 0's leaf nodes, we are guranteed to have leaf nodes with all 1's
  106. 	for(int i=0; i<n; i++){
  107. 		degree[i] = graph[i].size();
  108. 		if(degree[i] == 1 ){
  109. 			q.push(i);
  110. 		}
  111. 	}
  112.  
  113. 	while(!q.empty()){
  114. 		removed++;
  115. 		auto ch = q.front(); q.pop();
  116.  
  117. 		int parent = -1;
  118.  
  119. 		for(auto& par : graph[ch]){
  120. 			degree[par]--;
  121. 			parent = par;
  122. 		}
  123.  
  124. 		if(parent != -1){
  125. 			auto ch_to_delete = find(graph[parent].begin(), graph[parent].end(), ch);
  126. 			graph[parent].erase(ch_to_delete);
  127.  
  128. 			auto par_to_delete = find(graph[ch].begin(), graph[ch].end(), parent);
  129. 			graph[ch].erase(par_to_delete );
  130. 		}
  131. 	}
  132.  
  133. 	int left = n - removed;
  134.  
  135. 	if(left <= 0) cout << 0 << endl;
  136. 	else cout << 2 * (left - 1) << endl;
  137.  
  138. }
  139.  
  140.  
  141. void solve() {
  142.  
  143. 	int n;
  144. 	cin >> n;
  145.  
  146. 	coins.resize(n);
  147. 	edges.resize(n-1);
  148.  
  149. 	for(int i=0; i<n; i++) cin >> coins[i];
  150.  
  151. 	for(int i=0; i<n-1; i++){
  152. 		int x, y;
  153. 		cin >> x >> y;
  154. 		edges[i] = {x, y};
  155. 	}
  156.  
  157.     consistency(n);
  158.  
  159. }
  160.  
  161.  
  162.  
  163.  
  164.  
  165. int32_t main() {
  166.     ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
  167.  
  168.     int t = 1;
  169.     // cin >> t;
  170.     while (t--) {
  171.         solve();
  172.     }
  173.  
  174.     return 0;
  175. }
Success #stdin #stdout 0.01s 5288KB
comments (?)
stdin
copy 
11
0 0 1 1 0 0 0 0 0 1 1 
0 1
0 2
0 6
1 3
1 4
3 10
4 5
6 7
7 8 
8 9
stdout
copy 
6
https://ideone.com/lawUsD
language:
C++14 (gcc 8.3)
created:
9 hours 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!