CUBE-T165: National Ant Colony

#include <iostream>
#include <vector>
#include <tuple>
using namespace std;

#define tiii tuple<int, int, int>

vector<vector<int>> ants;
vector<tiii> edges; /// (u, v, w)
vector<int> DS;
vector<int> ranks;
int nv, ne, ns;

int DSFind(int x){
    if(DS[x] == x){
        return x;
    } else {
        return DS[x] = DSFind(DS[x]);
    }
}

void DSUnion(int x, int y){
    int headx = DSFind(x);
    int heady = DSFind(y);
    if(ranks[headx] > ranks[heady]){
        DS[heady] = headx;
    } else {
        DS[headx] = heady;
        if(ranks[headx] == ranks[heady]){
            ++ranks[heady];
        }
    }
    return;
}

int DestroyEdge(int limit){
    DS.assign(nv + 2, 0);
    ranks.assign(nv + 2, 1);
    for(int i = 1; i <= nv; ++i){
        DS[i] = i;
    }

    int cnt = 0;
    for(auto edge : edges){
        int u = get<0>(edge);
        int v = get<1>(edge);
        int w = get<2>(edge);
        if(w <= limit){
            ++cnt;
            continue;
        }
        if(DSFind(u) != DSFind(v)){
            DSUnion(u, v);
        }
    }

    for(int i = 1; i <= ns; ++i){
        int sz = ants[i].size();
        if(sz == 0 || sz == 1){
            continue;
        }
        for(int j = 0; j < sz - 1; ++j){
            if(DSFind(ants[i][j]) != DSFind(ants[i][j + 1])){
                return -1;
            }
        }
    }
    return cnt;
}

int main(){

    int x, w, u, v, mxlen, mnlen;

    scanf("%d %d %d", &nv, &ne, &ns);
    ants.assign(ns + 1, vector<int>());
    for(int i = 1; i <= nv; ++i){
        scanf("%d", &x);
        ants[x].push_back(i);
    }

    mxlen = 0;
    mnlen = 1e9;
    for(int i = 0; i < ne; ++i){
        scanf("%d %d %d", &u, &v, &w);
        edges.push_back(make_tuple(u, v, w));
        mxlen = max(mxlen, w);
        mnlen = min(mnlen, w);
    }

    int l = mnlen;
    int r = mxlen;
    int ans = 0;
    while(l <= r){
        int m = (l + r) / 2;
        int tmp = DestroyEdge(m);
        if(tmp == -1){
            r = m - 1;
        } else {
            ans = max(ans, tmp);
            l = m + 1;
        }
    }

    cout << ans;

    return 0;
}