National Ant Colony

#include<stdio.h>
#include<vector>
#include<algorithm>
#include<stack>
using namespace std;
struct edge{
    int v,w;
};

int main()
{
    int n,m,k;
    scanf("%d %d %d",&n,&m,&k);
    int Alledge[m];
    vector<edge> adj[n+1]; //adj list graph
    vector<int> ant[k+1]; //store pos of any ant type
    for(int i = 1 ; i <= n ; i ++) //input ant type by pos
    {
        int x;
        scanf("%d",&x);
        ant[x].push_back(i);
    }

    int maxw = 1; //max weight
    for(int i = 0 ; i < m ; i ++){//input edge
        int u,v,w;
        scanf("%d %d %d",&u,&v,&w);
        adj[u].push_back({v,w});
        adj[v].push_back({u,w});
        if(w > maxw)maxw = w;
        Alledge[i] = w;
    }
    sort(Alledge,Alledge+m);
    int p = 0;
    int l = 0 , r = maxw;
    while(l <= r){ //bsearch find best p
        int mid = (l+r)/2;

        // dfs write component variable
        int component[n+1];
        bool visited[n+1];
        for(int i = 0 ; i <= n ; i ++){
            component[i] = i;
            visited[i] = false;
        }

        //find node to start dfs write component
        for(int i = 1 ; i <= n ; i ++){
            if(!visited[i]){
                //dfs
                stack<int> st;
                st.push(i);
                while(!st.empty()){
                    int u = st.top();
                    st.pop();
                    if(visited[u])continue;
                    visited[u] = true;
                    component[u] = i;
                    //adj traversal
                    for(int j = 0 ; j < adj[u].size() ; j ++){
                        int v = adj[u][j].v;
                        int w = adj[u][j].w;
                        if(!visited[v] && w > mid){
                            st.push(v);
                        }
                    }
                }
            }
        }
        //checking ant
        int check = true;
        for(int i = 1 ; i <= k && check ; i ++){
            if(ant[i].size() <= 1)continue; //always true
            int com = component[ant[i][0]];
            for(int j = 1 ; j < ant[i].size() && check ; j ++){
                if(component[ant[i][j]] != com){
                    check = false;
                    break;
                }
            }
        }

        if(check){
            if(mid > p)p = mid;
            l = mid + 1;
        }else{
            r = mid - 1;
        }
    }
    l = 0 ; r = m - 1;
    int ans = 0;
    while(l <= r){
        int mid = (l+r)/2;
        if(Alledge[mid] <= p){
            ans = mid + 1;
            l = mid + 1;
        }
        else{
            r = mid - 1;
        }
    }
    printf("%d",ans);
    return 0;
}