Prim's MST.cpp

//Logic: https://www.hackerearth.com/practice/algorithms/graphs/minimum-spanning-tree/tutorial/
#include <bits/stdc++.h>
using namespace std;
int dist[N], parent[N];
bool vis[N];
vector<pair<int, int>> g[N], tree[N];

int primsMST(int source) //Finds the cost and makes the MST
{
	for (int i = 1; i <= n; i++)
		dist[i] = 1e18;
	set<pair<int, int>> s;
	s.insert({0, source});
	int cost = 0;
	dist[source] = 0;
	while (!s.empty())
	{
		auto x = *(s.begin());
		s.erase(x);
		vis[x.second] = 1;
		cost += x.first;
		int u = x.second;
		int v = parent[x.second];
		int w = x.first;
		tree[u].push_back({v, w});
		tree[v].push_back({u, w});
		for (auto it : g[x.second])
		{
			if (vis[it.first])
				continue;
			if (dist[it.first] > it.second)
			{
				s.erase({dist[it.first], it.first});
				dist[it.first] = it.second;
				s.insert({dist[it.first], it.first});
				parent[it.first] = x.second;
			}
		}
	}
	return cost;
}

//Sample Problem 1: https://codeforces.com/contest/609/problem/E
//Sample Solution 1: https://codeforces.com/contest/609/submission/39951860