fork download
copy 
  1. //NiceDuck
  2. #include "bits/stdc++.h"
  3. typedef long long ll;
  4. using namespace std;
  5. #define FILE "000"
  6. #define foru(i,a,b) for(int i=(int)(a); i<=(int)(b); ++i)
  7. #define ford(i,a,b) for(int i=(int)(a); i>=(int)(b); --i)
  8. #define fastio ios_base::sync_with_stdio(0);cin.tie(0);
  9. #define pb push_back
  10. #define fi first
  11. #define se second
  12. #define pii pair<int,int>
  13. #define pil pair<int,ll>
  14. #define pli pair<ll,int>
  15. #define MOD 1000000007
  16. #define el "\n"
  17.  
  18. const int MAX=2e5+5;
  19. int n,m;
  20. pair<int,int> q[MAX];
  21. ll ans[MAX];
  22. vector<int> adj[MAX];
  23. struct Edge 
  24. {
  25.     ll w; int u,v;
  26. };
  27. vector<Edge> edge;
  28. bool cmp(const Edge &x, const Edge &y)
  29. {
  30.     return x.w<y.w;
  31. }
  32.  
  33. int par[MAX],sz[MAX];
  34. void buildDsu()
  35. {
  36.     foru(i,1,n)
  37.     {
  38.         par[i]=i;
  39.         sz[i]=1;
  40.     }
  41. }
  42. int find_par(int u)
  43. {
  44.     if(par[u]==u) return u;
  45.     return par[u]=find_par(par[u]);
  46. }
  47. void join(int u, int v)
  48. {
  49.     u=find_par(u); v=find_par(v);
  50.     if(u==v) return;
  51.     if(sz[v]>sz[u]) swap(u,v);
  52.     par[v]=u;
  53.     sz[u]+=sz[v];
  54.     return;
  55. }
  56.  
  57. int main()
  58. {
  59.     fastio
  60.     if(fopen(FILE ".inp","r"))
  61.     {
  62.        freopen(FILE ".inp","r",stdin);
  63.        freopen(FILE ".out","w",stdout);
  64.     }
  65.  
  66.     cin>>n>>m; 
  67.     foru(i,1,n-1)
  68.     {
  69.         int u,v; ll w; cin>>u>>v>>w;
  70.         edge.pb({w,u,v});
  71.     }
  72.     buildDsu();
  73.     foru(i,1,m)
  74.     {
  75.         cin>>q[i].fi;
  76.         q[i].se=i;
  77.     }
  78.     sort(q+1,q+m+1);
  79.     if(edge.size()==0)
  80.     {
  81.         foru(i,1,m) cout<<"0 ";
  82.         return 0;
  83.     }
  84.     sort(edge.begin(),edge.end(),cmp);
  85.     int j=0;
  86.     ll curAns=0;
  87.     foru(i,1,m)
  88.     {
  89.         ans[q[i].se]=curAns;
  90.         if(edge[j].w>q[i].fi) continue;
  91.         else 
  92.         {
  93.             while(j<edge.size() && edge[j].w<=q[i].fi)
  94.             {
  95.                 int u=edge[j].u, v=edge[j].v;
  96.                 ++j;
  97.                 u=find_par(u); v=find_par(v);
  98.                 if(u==v) continue;
  99.                 curAns+=(1LL*sz[u]*sz[v]);
  100.                 join(u,v);
  101.             }
  102.         }
  103.         ans[q[i].se]=curAns;
  104.     }
  105.     foru(i,1,m) cout<<ans[i]<<' ';
  106.  
  107.     return 0;
  108. }
Success #stdin #stdout 0.01s 9736KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/jpHxGD
language:
C++14 (gcc 8.3)
created:
21 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!