[CF1213G] Path Queries - Kruskal

Posted 2021-03-03 mollnn

给你一个有 (n) 个点的带权树，有 (m) 个查询，每次查询最大权值不大于 (x) 的点对的数目

Solution

模拟 Kruskal 的过程，并将整个过程的答案记录下来

询问 (x) 的时候，二分找到最后一个 (le x) 的位置，输出这个时刻的答案即可

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 400005;

struct edge {
    int u,v,w;
    bool operator < (const edge &b) const {
        return w < b.w;
    }
} e[N];
int n,m,t1,t2,t3,ans,a[N],b[N],fa[N],sz[N];

int find(int p) {
    return fa[p]==p ? p : fa[p]=find(fa[p]);
}

void merge(int p,int q) {
    p=find(p); q=find(q);
    if(p-q) {
        fa[p]=q;
        ans+=sz[p]*sz[q];
        sz[q]+=sz[p];
    }
}

signed main() {
    ios::sync_with_stdio(false);
    cin>>n>>m;
    for(int i=1;i<n;i++) {
        cin>>t1>>t2>>t3;
        e[i]={t1,t2,t3};
    }
    sort(e+1,e+n);
    for(int i=1;i<=n;i++) {
        fa[i]=i;
        sz[i]=1;
    }
    for(int i=1;i<n;i++) {
        if(find(e[i].u)!=find(e[i].v)) {
            merge(e[i].u,e[i].v);
            a[i]=ans;
        }
        b[i]=e[i].w;
    }
    for(int i=1;i<=m;i++) {
        int x;
        cin>>x;
        cout<<a[upper_bound(b+1,b+n,x)-b-1]<<" ";
    }
}