fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. #define ll long long int
  4. const int M = 1e9 + 7;
  5.  
  6. int dfs(int node, vector<int> adj[], int color, vector<int> &colors, vector<int> &vis)
  7. {
  8.   if (colors[node - 1] != color)
  9.     return 0;
  10.   int sum = 1;
  11.   vis[node] = 1;
  12.   for (auto it : adj[node])
  13.   {
  14.     sum += dfs(it, adj, color, colors, vis);
  15.   }
  16.   return sum;
  17. }
  18.  
  19. int main()
  20. {
  21.   ios_base::sync_with_stdio(false);
  22.   cin.tie(NULL);
  23.   int n;
  24.   cin >> n;
  25.   vector<int> tree_from(n - 1);
  26.   vector<int> tree_to(n - 1);
  27.   vector<int> colors(n);
  28.  
  29.   for (int i = 0; i < n - 1; i++)
  30.     cin >> tree_from[i];
  31.   for (int i = 0; i < n - 1; i++)
  32.     cin >> tree_to[i];
  33.   for (int i = 0; i < n; i++)
  34.     cin >> colors[i];
  35.   vector<int> adj[n + 1];
  36.   for (int i = 0; i < n - 1; i++)
  37.   {
  38.     adj[tree_from[i]].push_back(tree_to[i]);
  39.   }
  40.   vector<int> vis(n + 1, 0);
  41.   ll ans = 0;
  42.   for (int i = 1; i <= n; i++)
  43.   {
  44.     if (vis[i] == 0)
  45.     {
  46.       int val = dfs(i, adj, colors[i - 1], colors, vis);
  47.       ans += max(0ll, ((val - 2) * 1ll * (val - 1)) / 2);
  48.     }
  49.   }
  50.  
  51.   cout << ans << endl;
  52.  
  53.   return 0;
  54. }
Success #stdin #stdout 0.01s 5320KB
comments (?)
stdin
copy 
6
1 2 2 1 5
2 3 4 5 6
0 1 1 1 0 0
stdout
copy 
2
https://ideone.com/CL1G9b
language:
C++14 (gcc 8.3)
created:
1 day 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!