fork download
copy 
  1. #include<bits/stdc++.h>
  2. using namespace std;
  3. #define endl '\n'
  4. #define int long long int
  5. const int MOD = 1000000007;
  6. const int MOD2 = 998244353;
  7. const int INF = LLONG_MAX / 2;
  8. const int MAXN = 100000;
  9. int primes[1000000];
  10.  
  11. /*void seive() {
  12.     fill(primes, primes + 1000000, 1);
  13.     primes[0] = primes[1] = 0;
  14.     for (int i = 2; i * i < 1000000; i++) {
  15.         if (primes[i]) {
  16.             for (int j = i * i; j < 1000000; j += i) {
  17.                 primes[j] = 0;
  18.             }
  19.         }
  20.     }
  21. }
  22.  
  23. bool isPrime(int n) {
  24.     if (n <= 1) return false;
  25.     for (int i = 2; i * i <= n; i++) {
  26.         if (n % i == 0) return false;
  27.     }
  28.     return true;
  29. }
  30.  
  31. int gcd(int a, int b) {
  32.     if (a == 0) return b;
  33.     return gcd(b % a, a);
  34. }*/
  35.  
  36. int power(int a, int b, int mod) {
  37.     int res = 1;
  38.     a %= mod;
  39.     while (b > 0) {
  40.         if (b & 1) res = res * a % mod;
  41.         a = a * a % mod;
  42.         b >>= 1;
  43.     }
  44.     return res;
  45. }
  46.  
  47. // nCr % MOD for n < MOD
  48. int nCrModP(int n, int r) {
  49.     if (r > n) return 0;
  50.     if (r == 0 || r == n) return 1;
  51.  
  52.     int numerator = 1, denominator = 1;
  53.     for (int i = 0; i < r; i++) {
  54.         numerator = (numerator * (n - i)) % MOD;
  55.         denominator = (denominator * (i + 1)) % MOD;
  56.     }
  57.     return (numerator * power(denominator, MOD - 2, MOD)) % MOD;
  58. }
  59.  
  60. // Lucas's Theorem
  61. int lucas(int n, int r) {
  62.     if (r == 0) return 1;
  63.     return (lucas(n / MOD, r / MOD) * nCrModP(n % MOD, r % MOD)) % MOD;
  64. }
  65. bool isPossible(int A[] , int n , int mid , int k){
  66.     int temp[n];
  67.     for(int i = 0; i < n; i++) {
  68.         temp[i] = A[i];
  69.     }
  70.     for(int i = 0 ; i < n ; i++){
  71.         if(temp[i]<=mid){
  72.            int d = mid-temp[i];
  73.            k -= d;
  74.         }
  75.     }
  76.     return k<=0;
  77. }
  78. void bfs(int node , vector<int>A[] , int visited[] , int &cnt){
  79.     queue<int>q;
  80.     q.push(node);
  81.     visited[node] = true;
  82.     cnt++;
  83.     while(!q.empty()){
  84.         int node = q.front();
  85.         q.pop();
  86.         for(auto node1 : A[node]){
  87.             if(!visited[node1]){
  88.                 cnt++;
  89.                 visited[node1] = true;
  90.                 q.push(node1);
  91.             }
  92.         }
  93.     }
  94. }
  95. void solve() {
  96.     int n;
  97.     cin>>n;
  98.     int B[n+1][2];
  99.     int colors[n+1];
  100.     vector<int>A[n+1];
  101.     int visited[n+1];
  102.     fill(visited,visited+n+1,false);
  103.     for(int i = 0 ; i<n-1 ; i++){
  104.         int u,v;
  105.         cin>>u>>v;
  106.         B[i][0] = u;
  107.         B[i][1] = v;
  108.     }
  109.     for(int i = 1 ; i<=n ; i++){
  110.         cin>>colors[i];
  111.     }
  112.     for(int i = 0 ; i<n-1 ; i++){
  113.         if(colors[B[i][0]]==colors[B[i][1]]){
  114.             A[B[i][0]].push_back(B[i][1]);
  115.             A[B[i][1]].push_back(B[i][0]);
  116.         }
  117.     }
  118.     int sum = 0;
  119.     for(int i = 1 ; i<=n ; i++){
  120.         if(!visited[i]){
  121.             int cnt = 0;
  122.             bfs(i,A,visited,cnt);
  123.             if(cnt>=3){
  124.             sum += ((cnt-1)*(cnt-2))/2;
  125.             }
  126.         }
  127.     }
  128.     cout<<sum<<endl;
  129. }
  130.  
  131. signed main() {
  132.     ios::sync_with_stdio(false); cin.tie(NULL);
  133.     //int t;
  134.     //cin >> t;
  135.     //while (t--) {
  136.         solve();
  137.     //}
  138.     return 0;
  139. }
  140.  
Success #stdin #stdout 0.01s 5320KB
comments (?)
stdin
copy 
6
1 2
2 3
2 4
1 5
5 6
0 1 1 1 0 0
stdout
copy 
2
https://ideone.com/tcSQUa
language:
C++14 (gcc 8.3)
created:
17 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!