RT 求助

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

int n,m,k,cntEdge,cntVer;

struct Edge {
    int u,v,w,use=0;
    Edge() {}
    Edge(int a,int b,int c) : u(a),v(b),w(c) {}
};

vector<Edge> edge;

class DSU {
    int fa[20010];
public :
    void init(int size) {
        for (int i=1; i<=size; ++i) {
            fa[i] = i;
        }
    }
    int find(int x) {
        if (x != fa[x]) {
            fa[x] = find(fa[x]);
        }
        return fa[x];
    }
    void unite(int x,int y) {
        fa[find(x)] = find(y);
    }
} dsu;

// 题意等价于，一个图存在n条权值为0或1的边，求边权和为k的生成树

void calculate() {
    dsu.init(n);
    sort(edge.begin(),edge.end(),[](Edge a,Edge b) {
        return a.w < b.w;  // 把0边放到前面
    });
    for (auto &ver : edge) {
        if (dsu.find(ver.u) != dsu.find(ver.v)) {
            dsu.unite(ver.u,ver.v);  // 0边全部加入
            if (ver.w == 1) { // 统计必加的1边个数（保证连通）
                ++cntEdge;
                ver.use = 1;
                ver.w = -1;
            }
        }
    }
}

void Kruskal() {
    dsu.init(n);
    sort(edge.begin(),edge.end(),[](Edge a,Edge b) {
        return a.w > b.w;  // 把1边放到前面
    });
    for (auto ver : edge) {  // 把必加的边全都加进去
        if (ver.use) {
            dsu.unite(ver.u,ver.v);
            ++cntVer;
        }
    }
    for (auto &ver : edge) {
        if (dsu.find(ver.u) != dsu.find(ver.v)) {
            if (ver.w == 1 and cntEdge < k) {
                dsu.unite(ver.u,ver.v);
                ++cntEdge, ++cntVer;
                ver.use = 1;
            } else if (ver.w == 0 and cntEdge == k) {
                dsu.unite(ver.u,ver.v);
                ++cntVer;
                ver.use = 1;
            }
        }
    }
}

int main() {
    scanf("%d %d %d",&n,&m,&k);
    for (int i=1; i<=m; ++i) {
        int u,v,w;
        scanf("%d %d %d",&u,&v,&w);
        if (w == 0) {  // 鹅卵石路，边权按1算
            edge.push_back({u,v,1});
        } else {       // 水泥路，边权按0算
            edge.push_back({u,v,0});
        }
    }

    // 先对0边（水泥路）跑一遍Kruskal，算出必须要加的1边的个数
    calculate();
    if (cntEdge > k) {  // 边权和至少为cntEdge，则无法达到k，否则不连通
        printf("No solution\n");
        return 0;
    }
    
    Kruskal();
    if (cntEdge < k or cntVer + 1 < n) {
        printf("No solution\n");
        return 0;
    }
    for (auto ver : edge) {
        if (ver.use) {
            printf("%d %d %d\n",ver.u,ver.v,1-(ver.w==-1?1:ver.w));
        }
    }
    return 0;
}