Prim
Algorithm
Implementation
// Time: O(V^2)
// Space: O(V)
int prim(int N, int M, vector<vector<int>> &edges) { // N nodes, M edges. edges is an array of `<u, v, w>` where u, v are from and to nodes, w is the weight.
vector<vector<pair<int, int>>> G(N);
for (auto &e : edges) {
int u = e[0] - 1, v = e[1] - 1, w = e[2];
G[u].emplace_back(v, w);
G[v].emplace_back(u, w);
}
int ans = 0, cur = 0; // start from a random node, say 0.
vector<int> dist(N, INT_MAX), seen(N);
for (int i = 0; i < N - 1; ++i) {
seen[cur] = 1; // add this node to the MST
for (auto &[v, w] : G[cur]) {
if (seen[v]) continue; // skip those neighbors that are already in the MST
dist[v] = min(dist[v], w); // Use the current node to update the distance of unvisited nodes.
}
cur = min_element(begin(dist), end(dist)) - begin(dist); // greedily pick an unconnected node with the minimum distance
ans += dist[cur];
dist[cur] = INT_MAX; // mark this distance as used
}
return ans;
}Reference
Last updated