fork download 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. typedef long long ll;
  4. #define Shoyo ios_base::sync_with_stdio(0);cin.tie(NULL);
  5. #define ff first
  6. #define ss second
  7. #define pii pair<ll,ll>
  8. #define all(v) v.begin(), v.end()
  9. #define allr(v) v.rbegin(), v.rend()
  10. #define el "\n"
  11. struct DSU {
  12.     vector<ll> parent;
  13.     vector<ll> sz;
  14.     ll comp=0;
  15.     void init(ll n) {
  16.         parent.assign(n,0);
  17.         sz.assign(n,1);
  18.         iota(all(parent),0);
  19.     }
  20.     DSU(ll n) {
  21.         init(n);
  22.     }
  23.     ll find(ll u) {
  24.         if (u==parent[u]) return u;
  25.         return parent[u]=find(parent[u]);
  26.     }
  27.     bool merge(ll u,ll v) {
  28.         u=find(u);
  29.         v=find(v);
  30.         if (u==v) return false;
  31.         if (sz[u]>sz[v])swap(u,v);
  32.         parent[u]=v;
  33.         sz[v]+=sz[u];
  34.         return true;
  35.     }
  36. };
  37.  
  38. void solve() {
  39.     ll n,q;cin>>n>>q;
  40.    vector<array<ll,3>>ad;
  41.     ad.reserve(n);
  42.    for(ll i=0;i<n-1;i++) {
  43.        ll u,v,w;cin>>u>>v>>w;
  44.        u--; v--;
  45.        ad.push_back({w,u,v});
  46.    }
  47.     vector<pii>v(q);
  48.     vector<ll>ans(q);
  49.     for(ll i=0;i<q;i++) {
  50.         cin>>v[i].ff;
  51.         v[i].ss=i;
  52.     }
  53.     sort(all(ad));
  54.     sort(all(v));
  55.     DSU dsu(n);
  56.     ll edge=0,j=0;
  57.     for(ll i=0;i<q;i++) {
  58.         ll x=v[i].ff;
  59.         while (j<n-1 and ad[j][0]<=x) {
  60.             auto [w,u,v]=ad[j];
  61.            edge+=(dsu.sz[u]*dsu.sz[v]);
  62.             dsu.merge(u,v);
  63.             j++;
  64.         }
  65.         ans[v[i].ss]=edge;
  66.     }
  67.  
  68.     for (ll i=0;i<q;i++) {
  69.         cout<<ans[i]<<" ";
  70.     }
  71.     cout<<endl;
  72. }
  73. int main(){
  74.     Shoyo;
  75.     ll t = 1;
  76.     // if(!(cin >> t)) return 0;
  77.     while (t--) {
  78.         solve();
  79.     }
  80.     return 0;
  81. }
Success #stdin #stdout 0.01s 5320KB
comments (?)
stdin 
Standard input is empty
stdout 

	

		
	 
	
	
			

					
			
			
			
			
			
			
			
			

			
			
			
		

	
		
		

			


			                        
			
			

				

	
					
					 https://ideone.com/FA6Uzr
					
					
					
					
					

												

												
						
						

					

					
					
					

						
language:

						
C++14 (gcc 8.3)

					

					

						
created:
								
						
3 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!