https://codeforces.com/problemset/problem/1213/G
把询问离线,对于每个询问,我们把比他的要求小的边全部加进树中。
那么答案就是所有连通块的贡献,每个连通块的贡献是sz*(sz-1)/2
那么我们按询问的要求值从小到大排序,然后每次尝试加边,并用并查集维护连通块大小
#include<bits/stdc++.h>
#define maxl 400010
using namespace std;
int n,m;
int a[maxl],ehead[maxl],f[maxl],t[maxl];
long long ans[maxl];
struct ed
{
int u,v,w;
}e[maxl];
struct que
{
int val,id;
}q[maxl];
char s[maxl];
inline bool cmp(const ed &x,const ed &y)
{
return x.w<y.w;
}
inline bool cmpq(const que &x,const que &y)
{
return x.val<y.val;
}
inline void prework()
{
scanf("%d%d",&n,&m);
int u,v,w;
for(int i=1;i<n;i++)
{
scanf("%d%d%d",&u,&v,&w);
e[i]=ed{u,v,w};
}
sort(e+1,e+1+n-1,cmp);
for(int i=1;i<=m;i++)
{
scanf("%d",&q[i].val);
q[i].id=i;
}
sort(q+1,q+1+m,cmpq);
}
inline int find(int x)
{
if(f[x]!=x)
f[x]=find(f[x]);
return f[x];
}
inline void mainwork()
{
for(int i=1;i<=n;i++)
{
f[i]=i;
t[i]=1;
}
long long tmp=0;
int id=1,x,y;
for(int i=1;i<=m;i++)
{
while(e[id].w<=q[i].val && id<=n-1)
{
x=find(e[id].u);y=find(e[id].v);
if(x!=y)
{
tmp-=1ll*(t[x]-1)*t[x]/2;
tmp-=1ll*(t[y]-1)*t[y]/2;
f[y]=x;t[x]+=t[y];
tmp+=1ll*(t[x]-1)*t[x]/2;
}
id++;
}
ans[q[i].id]=tmp;
}
}
inline void print()
{
for(int i=1;i<=m;i++)
printf("%lld%c",ans[i],(i==m)?'\n':' ');
}
int main()
{
prework();
mainwork();
print();
return 0;
}