fork download
copy 
  1. #include<bits/stdc++.h>
  2. using namespace std;
  3. using ll = long long;
  4. using ld = long double;
  5. #define endme "\n";
  6. #define ffprint(x)        \
  7.     for (auto &i : x)      \
  8.     {                     \
  9.         cout << i << " "; \
  10.     } \
  11.     cout<<endme;
  12. #define vectinput(x) for(auto &i : x) cin >> i;
  13. #ifndef ONLINE_JUDGE
  14.    #include"Debug.cpp"
  15. #endif
  16.  
  17. class DSUrank{
  18.     public:
  19.     vector<int>parent;
  20.     vector<int>rank;
  21.     DSUrank(int n){
  22.         parent.resize(n);
  23.         rank.resize(n);
  24.         for(int i=0;i<n;i++)parent[i]=i;
  25.     }
  26.     int find_parent(int x){
  27.         if(parent[x]==x)return x;
  28.         parent[x]=find_parent(parent[x]);
  29.         return parent[x];
  30.     }
  31.     void union_(int x,int y){
  32.         int pary=find_parent(y);
  33.         int parx=find_parent(x);
  34.         if(pary==parx)return;
  35.         if(rank[parx]>rank[pary]){
  36.             parent[pary]=parx;
  37.         }else if(rank[parx]<rank[pary]){
  38.             parent[parx]=pary;
  39.         }else{
  40.             parent[parx]=pary;
  41.             rank[pary]++;
  42.         }
  43.     }
  44. };
  45.  
  46. void solve(){
  47.     int n;
  48.     cin>>n;
  49.     vector<pair<int,int>>edges(n-1);
  50.     vector<int>color(n);
  51.     for(auto &i:edges){
  52.         cin>>i.first>>i.second;
  53.         i.first--;
  54.         i.second--;
  55.     }
  56.     vectinput(color);
  57.  
  58.     DSUrank mydsu(n);
  59.     for(auto &i:edges){
  60.         if(color[i.first]==color[i.second]){
  61.             mydsu.union_(i.first,i.second);
  62.         }
  63.     }
  64.     map<int,int> mp;
  65.     for(int i=0;i<n;i++){
  66.         int x=mydsu.find_parent(i);
  67.         mp[x]++;
  68.     }
  69.     int answer=0;
  70.     for(auto &i:mp){
  71.         answer+=((i.second-1)*(i.second-2))/2;
  72.     }
  73.     cout<<answer<<endme;
  74.  
  75. }
  76.  
  77.  
  78. int main() {
  79.     ios::sync_with_stdio(false);
  80.     cin.tie(0);
  81.     cout.tie(0);
  82.     #ifndef ONLINE_JUDGE
  83.         freopen("0-InputFile/input.txt", "r", stdin);
  84.         freopen("0-InputFile/output.txt", "w", stdout);
  85.     #endif
  86.  
  87.    int t;
  88.    cin>>t;
  89. //    t=1;
  90.     while (t--) {
  91.         solve();
  92.     }
  93.     return 0;
  94. }
  95.  
  96. /*
  97. reset_template.bat
  98. %%%%%%%%+++*%@@@@@%%@%%%%%%%%%%%@@%%%%%%##**###%%%
  99. %%%%%%%%#*#%@@@@@@@%%%%%#%%%@@@@@%####**####%%@@@@
  100. %%%%%%%%##%@@@@@@@*+++++++++*#@@@@%####%%%%@@@@@@@
  101. %%%%%%%%#%%@@@@@#++++++++++++@@@@@@@%%%@@@@@@@@@@@
  102. @%#%%#%#%%%%%%*++++++++++++++*@@@@@@#++++++*######
  103. %%%%%#%%%###++++++++++++++++++*#@@@#*******#******
  104. #%%%%%%%%%%*+++++******+++++++++**++*##**+++***##%
  105. %%#%%%%%@@*++++*%@@#*****#%#*++++++++++++**##%@@@%
  106. %%%%%%%%%#+***#@@@#++*++*@@@#***+++++++*#%@@%##%%%
  107. %%%%%%%%#**###%@@#+++++++#@@@#*#**++**##%#####%@@@
  108. %%%%%%%%%**####%#++*##*++*@@@%###***#@@@@@%%%%@@@@
  109. %%%%%%%%%#*#####**#@@@@@***%%%%%##*%@@@@@@@@%%@@%%
  110. %%%%%%%%#%**######%%@@%##%%%%%%%##%@@@@@@@@@%#****
  111. %@@%%%###@@%#%%%@%%@@@@@@@%%%%%%%@@@@@@@@@@@@%#**#
  112. %%%%%%%%%@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@###
  113. %%%#%%%%%@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
  114. %%%%%%%%%@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@%
  115. %%%%%%%%%@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@%
  116. */
  117.  
  118.  
  119.  
  120.  
Success #stdin #stdout 0s 5324KB
comments (?)
stdin
copy 
1
6
1 2
2 3
2 4
1 5
5 6
0 1 1 1 0 0
stdout
copy 
2
https://ideone.com/AJoWJc
language:
C++ (gcc 8.3)
created:
19 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!