BZOJ 3624 免费道路

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define maxv 20050
#define maxe 100500
using namespace std;
int n,m,k,x,y,z,cnt0=0,cnt1=0,fath[maxv];
struct edge
{
    int u,v;
}e0[maxe],e1[maxe];
bool vis[maxe],vis0[maxe],vis1[maxe];
void reset()
{
    for (int i=1;i<=n;i++)
        fath[i]=i;
}
int getfather(int x)
{
    if (x!=fath[x])
        fath[x]=getfather(fath[x]);
    return fath[x];
}
void focus()
{
    reset();
    for (int i=1;i<=cnt1;i++)
    {
        int u=e1[i].u,v=e1[i].v;
        int f1=getfather(u),f2=getfather(v);
        if (f1!=f2) fath[f1]=f2;
    }
    for (int i=1;i<=cnt0;i++)
    {
        int u=e0[i].u,v=e0[i].v;
        int f1=getfather(u),f2=getfather(v);
        if (f1!=f2) vis[i]=true;
    }
}
bool kruskal()
{
    reset();
    int ret=0;
    for (int i=1;i<=cnt0;i++)
    {
        if (vis[i])
        {
            int u=e0[i].u,v=e0[i].v;
            int f1=getfather(u),f2=getfather(v);
            if ((f1!=f2) && (ret<k))
            {
                ret++;vis0[i]=true;
                fath[f1]=f2;
            }
        }
    }
    for (int i=1;i<=cnt0;i++)
    {
        if (!vis[i])
        {
            int u=e0[i].u,v=e0[i].v;
            int f1=getfather(u),f2=getfather(v);
            if ((f1!=f2) && (ret<k))
            {
                ret++;vis0[i]=true;
                fath[f1]=f2;
            }
        }
    }
    if (ret!=k) return false;
    for (int i=1;i<=cnt1;i++)
    {
        int u=e1[i].u,v=e1[i].v;
        int f1=getfather(u),f2=getfather(v);
        if (f1!=f2)
        {
            fath[f1]=f2;
            ret++;vis1[i]=true;
        }
    }
    if (ret!=n-1) return false;
    return true;
}
void print_e()
{
    for (int i=1;i<=cnt0;i++)
    {
        if (vis0[i])
            printf("%d %d 0\n",e0[i].u,e0[i].v);
    }
    for (int i=1;i<=cnt1;i++)
    {
        if (vis1[i])
            printf("%d %d 1\n",e1[i].u,e1[i].v);
    }
}
int main()
{
    scanf("%d%d%d",&n,&m,&k);
    for (int i=1;i<=m;i++)
    {
        scanf("%d%d%d",&x,&y,&z);
        if (!z)
        {
            cnt0++;
            e0[cnt0].u=x;e0[cnt0].v=y;
        }
        else
        {
            cnt1++;
            e1[cnt1].u=x;e1[cnt1].v=y;
        }
    }
    focus();
    if (!kruskal()) printf("no solution\n");
    else print_e();
    return 0;
}