TIOJ 1795 . 咕嚕咕嚕呱啦呱啦

給個簡單圖，邊權是0或1，問存不存在權重和k的生成樹。

#include <cstdio>
#include <cstdio>
#include <vector>

using namespace std;

struct pr{int x, y; pr(int _x = 0, int _y = 0):x(_x), y(_y){return;}};

struct djs{
    int ft[100000];
    djs(){
        for(int lx = 0;lx < 100000;lx++)
            ft[lx] = lx;
        return;
    }
    int query(int a){
        return ft[a] == a ? a : ft[a] = query(ft[a]);
    }
    void join(int a, int b){
        int fa = query(a), fb = query(b);
        ft[fa] = fb;
        return;
    }
};

djs dj0, dj1;

vector<pr> eg[2];

int main(){
    int n, m, k; scanf("%d %d %d", &n, &m, &k);
    for(int lx = 0;lx < m;lx++){
        int a, b, c; scanf("%d %d %d", &a, &b, &c);
        eg[c].push_back(pr(a-1, b-1));
    }
   
    for(pr aa : eg[0]) dj0.join(aa.x, aa.y);
   
    int prc = 0;
    for(pr aa : eg[1]){
        if(dj0.query(aa.x) == dj0.query(aa.y)) continue;
        dj0.join(aa.x, aa.y);
        dj1.join(aa.x, aa.y);
        prc++;
    }
   
    bool con = true;
    for(int lx = 1;lx < n and con;lx++)
        if(dj0.query(lx) != dj0.query(0))
            con = false;
    if(!con){
        puts("NIE");
        return 0;
    }

    for(pr aa : eg[1]){
        if(prc == k) break;
        if(dj1.query(aa.x) == dj1.query(aa.y)) continue;
        dj1.join(aa.x, aa.y);
        prc++;
    }

    puts(prc == k ? "TAK" : "NIE");
    return 0;
}