Prim's MST : Special Subtree Solution

#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>
#include <utility>
#include <bitset>
using namespace std;
auto comp=[](const pair<int,int> &l,const pair<int,int> &r){
    return l>r;
};
vector<vector<pair<int,int> > > a(3001);
priority_queue<pair<int,int>,vector<pair<int,int> >,decltype(comp) > pq(comp);
bitset<3001> taken;
void process(int u){
    taken[u]=1;
    for(int j=0;j<a[u].size();j++){
        pair<int,int> v=a[u][j];
        if(!taken[v.second])
            pq.push({v.first,v.second});
    }
}
int main(){
    int n,e;
    cin>>n>>e;
    for(int i=0;i<e;i++){
        int u,v,w;
        cin>>u>>v>>w;
        a[u].push_back({w,v});
        a[v].push_back({w,u});
    }
    int s,ans=0;
    cin>>s;
    process(s);
    while(!pq.empty()){
        pair<int,int> u=pq.top();
        pq.pop();
        int v=u.second;
        int w=u.first;
        if(!taken[v])
            {
                ans+=w;

                // cout<<ans<<endl;
                process(v);
            }
    }
    cout<<ans<<endl;
}

Comments