离线+生成树+并查集——cf1213G

#include<bits/stdc++.h>
using namespace std;
#define N 200005
#define ll long long 
struct Edge{int u,v,w;}e[N];
struct Query{int id,qq;}qq[N];
int n,q;
ll ans[N],sum;
bool operator<(Edge a,Edge b){return a.w<b.w;}
bool operator<(Query a,Query b){return a.qq<b.qq;}

int F[N],size[N];
int find(int x){
    if(F[x]==x)return x;
    int f=find(F[x]);
    return F[x]=f;
}
void bing(int u,int v){
    int f1=find(u),f2=find(v);
    if(f1==f2)return;
    sum-=(ll)size[f1]*(size[f1]-1)/2;
    sum-=(ll)size[f2]*(size[f2]-1)/2;
    size[f1]+=size[f2];
    sum+=(ll)size[f1]*(size[f1]-1)/2;
    F[f2]=f1;
}

int main(){
    cin>>n>>q;
    for(int i=1;i<n;i++)
        scanf("%d%d%lld",&e[i].u,&e[i].v,&e[i].w);        
    for(int i=1;i<=q;i++)
        scanf("%d",&qq[i].qq),qq[i].id=i;
    sort(e+1,e+n);
    sort(qq+1,qq+1+q);
    
    for(int i=1;i<=n;i++)F[i]=i,size[i]=1;
    
    int p=1;
    for(int i=1;i<=q;i++){
        while(qq[i].qq>=e[p].w && p<=n-1){
            int u=e[p].u,v=e[p].v;
            bing(u,v);
            p++;
        }
        ans[qq[i].id]=sum;
    }
    for(int i=1;i<=q;i++)
        cout<<ans[i]<<" ";
}