两次生成树

Sample Input
5 7 2
1 3 0
4 5 1
3 2 0
5 3 1
4 3 0
1 2 1
4 2 1

Sample Output
3 2 0
4 3 0
5 3 1
1 2 1

第一次构建生成树确定哪些鹅卵石路是必须免费的
第二次构建生成树是把生成树中除去必须免费的鹅卵石路的其他的路标记出来

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int maxn = 100005;
int pre[maxn],n,m,k;

struct node
{
    int u,v,w,flag;
}a[maxn];

int fnd(int x)
{
    if(x != pre[x])
    {
        pre[x] = fnd(pre[x]);
    }
    return pre[x];
}

int main()
{
    scanf("%d%d%d",&n,&m,&k);
    int num = 0;
    for(int i=1; i<=m; i++)
    {
        scanf("%d%d%d",&a[i].u,&a[i].v,&a[i].w);
        if(!a[i].w)
            num++;
    }
    if(num < k)
        return 0;
    for(int i=1; i<=n; i++)
    {
        pre[i] = i;
    }
    for(int i=1; i<=m; i++)//第一次构建生成树
    {
        int u = a[i].u;
        int v = a[i].v;
        int w = a[i].w;
        if(fnd(u) != fnd(v) && w)//先尽量让水泥路免费
        {
            pre[fnd(v)] = fnd(u);
        }
    }
    int cnt = 0;
    for(int i=1; i<=m; i++)
    {
        int u = a[i].u;
        int v = a[i].v;
        int w = a[i].w;
        if(fnd(u) != fnd(v) && !w)//统计哪些鹅卵石路是必须免费的
        {
            pre[fnd(v)] = fnd(u);
            a[i].flag = 1;
            cnt++;
        }
    }
    if(cnt > k)//必须免费的鹅卵石路的数量大于k，no solution
    {
        printf("no solution\n");
        return 0;
    }

    int sum = 0;
    for(int i=1; i<=n; i++)
    {
        pre[i] = i;
    }
    for(int i=1; i<=m; i++)//第二次构建生成树
    {
        int u = a[i].u;
        int v = a[i].v;
        int w = a[i].w;
        if(a[i].flag)//先把必须免费的鹅卵石放入生成树
        {
            pre[fnd(v)] = fnd(u);
            sum++;
        }
        else if(fnd(u) != fnd(v) && !w && cnt < k)//如果免费的鹅卵石路还没达到k，则再把这些鹅卵石路变为免费
        {
            pre[fnd(v)] = fnd(u);
            a[i].flag = 1;
            cnt++;
            sum++;
        }
    }
    if(cnt != k)//所有可以加入生成树的鹅卵石路数量少于k
    {
        printf("no solution");
        return 0;
    }
    for(int i=1; i<=m; i++)//把生成树中免费的水泥路标记出来
    {
        int u = a[i].u;
        int v = a[i].v;
        int w = a[i].w;
        if(fnd(u) != fnd(v) && w)
        {
            pre[fnd(v)] = fnd(u);
            a[i].flag = 1;
            sum++;
        }
    }
    if(sum != n-1)//无法构成生成树
    {
        printf("no solution");
        return 0;
    }
    for(int i=1; i<=m; i++)
    {
        if(a[i].flag)
            printf("%d %d %d\n",a[i].u,a[i].v,a[i].w);
    }
}