fork download
copy 
  1. #include<iostream>
  2. #include <ext/pb_ds/assoc_container.hpp>
  3. #include <ext/pb_ds/tree_policy.hpp>
  4. #define ll long long 
  5. #define pb push_back
  6. #define  fbo(x) find_by_order(x)
  7. #define  ook(x) order_of_key(x)
  8. using namespace std;
  9. using namespace __gnu_pbds;
  10. typedef tree<
  11. int,
  12. null_type,
  13. less<int>,
  14. rb_tree_tag,
  15. tree_order_statistics_node_update>
  16. ordered_set;
  17.  
  18. ll ans=0;
  19. int  checks(int n,vector<vector<int>>&adj)
  20. {
  21.        for(int i=1;i<=n;i++)
  22.        {
  23.        	 if(adj[i].size()==n-1)
  24.        	 	return i;
  25.        }
  26.        return -1;
  27. }
  28.  
  29. void dfss(int c,vector<vector<int>>&adj,int p2)
  30. {
  31. 	ll sm=0,bg=0,sm1=0,bg1=0;
  32. 	for(auto &i:adj[c])
  33. 	{
  34. 		if(i<c)
  35. 		{
  36. 			sm+=sm1;
  37. 			sm1++;
  38. 		}
  39. 		else if(i>c)
  40. 		{
  41. 			bg+=bg1;
  42. 			bg1++;
  43. 		}
  44. 	}
  45. 	if(p2==1)
  46. 		ans=bg;
  47. 	else if(p2==3)
  48. 		ans=sm;
  49. 	else if(p2==2)
  50. 		ans=sm1*bg1;
  51. 	return;
  52.  
  53. }
  54. void dfs(int u,int p,ordered_set &s,vector<vector<int>>&adj,int n,vector<bool>&vis,int p2)
  55. {
  56. 	  ll ts=u-1;
  57. 	  ll tb=n-u;
  58. 	  s.insert(u);
  59. 	  ll ps=s.ook(u);
  60. 	  ll pb=(ll)s.size()-s.ook(u)-1;
  61. 	  ll es=0;
  62. 	  ll eb=0;
  63. 	  ll ans1=0;
  64. 	  for(auto&v:adj[u])
  65. 	  {
  66. 	  	if(!vis[v])
  67. 	  	{
  68. 	  		vis[v]=true;
  69. 	  		dfs(v,u,s,adj,n,vis,p2);
  70.             ll cs=s.ook(u);
  71.             ll cb=(ll)s.size()-s.ook(u)-1;
  72.             ll ms=max(cs-ps,0LL);
  73.             ll mb=max(cb-pb,0LL);
  74.             if(p2==1)
  75.             ans1+=(mb)*(tb-mb);
  76.             else if(p2==3)
  77.             ans1+=(ms)*(ts-ms);
  78.             else if(p2==2)
  79.             ans1+=(ms)*(tb-mb)+mb*(ts-ms); 
  80.  
  81. 	  	ps=cs;
  82. 	  	pb=cb;
  83. 	  	es+=ms;
  84. 	  	eb+=mb;
  85. 	  }
  86. 	  }
  87. 	  ll sm=ts-es;
  88. 	  ll bg=tb-eb;
  89.       sm=max(0LL,sm);
  90.       bg=max(0LL,bg);
  91.       if(p2==1)
  92.       	ans1+=(bg)*(tb-bg);
  93.       else if(p2==3)
  94.       	ans1+=(sm)*(ts-sm);
  95.       else if(p2==2)
  96.       {
  97.       	ans1+=((sm)*(tb-bg)+bg*(ts-sm));
  98.       }
  99.       ans1/=2;
  100.       ans+=ans1;
  101.  
  102. }
  103. int main()
  104. {    ios::sync_with_stdio(0);
  105.         cin.tie(0);
  106.         cout.tie(0);
  107. 	int t;
  108. 	cin>>t;
  109. 	while(t--)
  110. 	{
  111. 		int n,p1,p2,p3,u,v;
  112. 		cin>>n>>p1>>p2>>p3;
  113. 		vector<vector<int>>adj(n+1);
  114. 		for(int i=1;i<n;i++)
  115. 		{
  116. 			cin>>u>>v;
  117. 			adj[u].pb(v);
  118. 			adj[v].pb(u);
  119. 		}
  120. 		ordered_set s;
  121.         ans=0;
  122. 	   int c=checks(n,adj);
  123. 	   if(c!=-1)
  124. 	   {
  125. 	   	  dfss(c,adj,p2);
  126. 	   	  cout<<ans<<"\n";
  127. 	   	  continue;
  128. 	   }
  129.         vector<bool>vis(n+1,0);
  130.         vis[1]=true;
  131.         s.insert(1);
  132.         dfs(1,0,s,adj,n,vis,p2);
  133.       //  cout<<"sasqs";
  134.        cout<<ans<<"\n";
  135.  
  136.  
  137. 	}
  138. 	return 0;
  139.  
  140. }
Success #stdin #stdout 0s 4492KB
comments ()
stdin
copy 
1
4
2 1 3
1 2
2 3
2 4
stdout
copy 
1
https://ideone.com/4RNFSV
language:
C++14 (gcc 8.3)
created:
5 years 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!