bzoj 3624: [Apio2008]免费道路 生成树的构造

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>
using namespace std;
#define MAXN 1000010
struct edge
{
        int x,y;
};
vector<edge> v0,v1;
bool vis0[MAXN],vis1[MAXN];
int uf[MAXN];
int get_fa(int now)
{
        return (uf[now]==now)?now:uf[now]=get_fa(uf[now]);
}
bool comb(int x,int y)
{
        x=get_fa(x);
        y=get_fa(y);
        if (x==y)return 0;
        uf[x]=y;
        return 1;
}
int main()
{
        freopen("input.txt","r",stdin);
        int n,m,i,j,k,x,y,z;
        int cnt=0;;
        int tot=0;
        scanf("%d%d%d",&n,&m,&k);
        edge et;
        for (i=0;i<=n;i++)
                uf[i]=i;
        for (i=0;i<m;i++)
        {
                scanf("%d%d%d",&et.x,&et.y,&z);
                if (z)
                        v1.push_back(et);
                else
                        v0.push_back(et);
        }
        for (i=0;i<v1.size();i++)
        {
                comb(v1[i].x,v1[i].y);
        }
        for (i=0;i<v0.size();i++)
        {
                vis0[i]=comb(v0[i].x,v0[i].y);
                cnt+=vis0[i];
        }
        if (cnt>k)
        {
                printf("no sulotion\n");
                return 0;
        }
        for (i=0;i<=n;i++)
                uf[i]=i;
        for (i=0;i<v0.size();i++)
        {
                if (vis0[i])
                        comb(v0[i].x,v0[i].y);
        }
        for (i=0;i<v0.size();i++)
        {
                if (cnt==k)break;
                if (vis0[i])continue;
                vis0[i]=comb(v0[i].x,v0[i].y);
                cnt+=vis0[i];
        }
        for (i=0;i<v1.size();i++)
        {
                vis1[i]=comb(v1[i].x,v1[i].y);
        }
        for (i=0;i<v0.size();i++)
                if (vis0[i])
                        tot++;
        for (i=0;i<v1.size();i++)
                if (vis1[i])
                        tot++;
        if (cnt!=k || tot!=n-1)
        {
                printf("no solution\n");
                return 0;
        }
        for (i=0;i<v0.size();i++)
                if (vis0[i])
                        printf("%d %d %d\n",v0[i].x,v0[i].y,0);
        for (i=0;i<v1.size();i++)
                if (vis1[i])
                        printf("%d %d %d\n",v1[i].x,v1[i].y,1);
}