// for practicing at https://www.hackerrank.com/challenges/primsmstsub/problem?isFullScreen=false

void addPair(unsigned, unsigned);
bool findLoop(unsigned startNode, unsigned numVertices);
void removePair(unsigned a, unsigned b);
bool pairExists(unsigned a, unsigned b);

vector <set<int>> Adj;      // Adj list(set) representation of MST

// Complete the prims function below.

int prims(int n, vector<vector<int>> edges, int start) {

    // save width as first element to keep sorted by width
    multimap<int, pair<int,int>> edgesSortedByWeights;      // can use priority queue too
    for(auto it=edges.begin(); it!=edges.end();it++)
    {
        edgesSortedByWeights.insert(pair<int, pair<int,int>>((*it)[2], pair<int,int>((*it)[0], (*it)[1])));
    }
    unsigned MSTWeight = 0;    // weight of the MST
    Adj.resize(n+1);        // 1 to n

    int Ai, Bi;
    auto it = edgesSortedByWeights.begin();
    MSTWeight = it->first;
    pair<unsigned, unsigned> AiBi = it->second;
    Ai = AiBi.first;
    Bi = AiBi.second;

    unsigned rootnode = Ai;     // let us use Ai of the smallest edge weight as first node
    addPair(Ai, Bi);            // first edge, has to go in
    it++;
    for (; it != edgesSortedByWeights.end(); it++)
    {
        AiBi = it->second;
        Ai = AiBi.first;
        Bi = AiBi.second;
        if(!pairExists(Ai,Bi))
        {
            addPair(Ai, Bi);
            if (findLoop(rootnode, n))
                removePair(Ai, Bi);
            else
            {
                MSTWeight += it->first;
            }
        }
    }
    //cout << MSTWeight<<endl;
    return MSTWeight;
}

void addPair(unsigned a, unsigned b)
{
    Adj[a].insert(b);
    Adj[b].insert(a);
}

void removePair(unsigned a, unsigned b)
{
    Adj[a].erase(b);
    Adj[b].erase(a);
}

bool dfsCycle(unsigned node, unsigned parent, vector<bool>&visited)
{
    bool result = false;
    if (!visited[node])
    {
        visited[node] = true;
        for (auto it = Adj[node].begin(); it != Adj[node].end() && !result; it++)
        {
            if (!visited[*it])
                result = dfsCycle(*it, node, visited);
            else
            {
                result = *it != parent;
            }
        }
    }
    return result;
}

bool findLoop(unsigned startNode, unsigned numVertices)
{
    vector<bool>visited(numVertices+1);
    unsigned parent = 0;
    return dfsCycle(startNode, parent, visited);
}

bool pairExists(unsigned a, unsigned b)
{
    return (Adj[a].find(b) != Adj[a].end()) && (Adj[b].find(a) != Adj[b].end());
}