BZOJ 3624: [Apio2008]免费道路 [生成树 并查集]

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int N=2e4+5, M=1e5+5;
typedef long long ll;
inline int read(){
    char c=getchar();int x=0,f=1;
    while(c<‘0‘||c>‘9‘){if(c==‘-‘)f=-1;c=getchar();}
    while(c>=‘0‘&&c<=‘9‘){x=x*10+c-‘0‘;c=getchar();}
    return x*f;
}

int n, m, k, u, v, c, m0, m1, p;
struct meow{int u, v, c;}a[M], b[M], ans[N];
int fa[N];
int find(int x) {return x==fa[x] ? x : fa[x]=find(fa[x]);}
int flag[N];
int main() {
    freopen("in","r",stdin);
    n=read(); m=read(); k=read();
    for(int i=1; i<=m; i++) {
        u=read(), v=read(), c=read();
        if(c==1) a[++m1]=(meow){u,v,c};
        else b[++m0]=(meow){u,v,c};
    }

    int cnt=0;
    for(int i=1; i<=n; i++) fa[i]=i;
    for(int i=1; i<=m1; i++) {
        u=a[i].u, v=a[i].v;
        int x=find(u), y=find(v);
        if(x==y) continue;
        fa[x]=y;
        if(++cnt == n-1) break;
    }
    for(int i=1; i<=m0; i++) {
        u=b[i].u, v=b[i].v;
        int x=find(u), y=find(v);
        if(x==y) continue;
        fa[x]=y; ans[++p]=b[i]; flag[i]=1;
        if(++cnt == n-1) break;
    }
    if(p > k || cnt < n-1) {puts("no solution"); return 0;}

    for(int i=1; i<=n; i++) fa[i]=i;
    for(int i=1; i<=p; i++) fa[find(ans[i].u)] = find(ans[i].v);
    cnt=p;
    if(cnt<k) for(int i=1; i<=m0; i++) if(!flag[i]){
        u=b[i].u, v=b[i].v;
        int x=find(u), y=find(v);
        if(x==y) continue;
        fa[x]=y; ans[++cnt]=b[i];
        if(cnt == k) break;
    }
    for(int i=1; i<=m1; i++) {
        u=a[i].u, v=a[i].v;
        int x=find(u), y=find(v);
        if(x==y) continue;
        fa[x]=y; ans[++cnt]=a[i];
        if(cnt == n-1) break;
    }
    for(int i=1; i<=cnt; i++) printf("%d %d %d\n",ans[i].u, ans[i].v, ans[i].c);
}