fork(1) download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std; 
  3. #define ll long long
  4. #define fi first
  5. #define se second
  6. #define int long long
  7. typedef pair<int,int> pii;
  8. const int MAXN = 3e5 +5;
  9. int n, ans, c, r, b;
  10. vector <vector<int>> adj;
  11. vector <int> col, red, blue;
  12. vector <bool> vis;
  13.  
  14. void dfs(int x){
  15.     vis[x] = true;
  16.     if (col[x] == 1)red[x]++;
  17.     if ( col[x] == 2) blue[x]++;
  18.     for (auto i : adj[x]){
  19.         if (vis[i]) continue;
  20.         dfs(i);
  21.         red[x] += red[i];
  22.         blue[x] += blue[i];
  23.     }
  24.     if ((red[x] == r && !blue[x]) || (blue[x] == b && !red[x])) ans++;
  25. }
  26.  
  27. signed main(){
  28.  
  29.     cin >> n;
  30. vis.resize(n+1);
  31. adj.resize(n+1);
  32. col.resize(n+1);
  33. red.resize(n+1);
  34. blue.resize(n+1);
  35.     for (int i =1; i <= n; i++){
  36.         cin >> c;
  37.         if ( c == 1) r++;
  38.         else if ( c== 2) b++;
  39.         col[i] = c;
  40.     }
  41.     for (int i =1; i <= n-1; i++){
  42.         int x, y;
  43.         cin >> x >> y;
  44.         adj[x].push_back(y);
  45.         adj[y].push_back(x);
  46.     }
  47.  
  48.     dfs(1);
  49.     cout << ans;
  50.  
  51. }
Success #stdin #stdout 0.01s 5288KB
comments (?)
stdin
copy 
3
1 1 2
2 3
1 3
stdout
copy 
Standard output is empty
https://ideone.com/tHs1GI
language:
C++ (gcc 8.3)
created:
10 days 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!