An intuitive Prim algorithm impl.
#include <vector> #include <iostream> #include <queue> #include <unordered_map> #include <unordered_set> using namespace std; struct Edge { Edge() :s(0), t(0), d(0) {} Edge(unsigned rs, unsigned rt, unsigned rd) : s(rs), t(rt), d(rd) {} unsigned s; unsigned t; unsigned d; bool operator()(const Edge &e1, const Edge &e2) { return e1.d > e2.d; } }; int main() { long n, m; cin >> n >> m; // from -> to -> (length, id) unordered_map<unsigned, unordered_map<unsigned, unsigned>> g; for (int i = 0; i < m; i++) { unsigned a, b, d; cin >> a >> b >> d; g[a][b] = g[b][a] = d; } unsigned s; cin >> s; unordered_set<unsigned> pickedSet; // linked node inx priority_queue<Edge, vector<Edge>, Edge> q; for (auto &kv : g[s]) { q.push(Edge(s, kv.first, kv.second)); } pickedSet.insert(s); unsigned long long ret = 0; while (!q.empty()) { Edge picked = q.top(); q.pop(); if (pickedSet.count(picked.t)) continue; pickedSet.insert(picked.t); // add new choices for (auto &kv : g[picked.t]) { q.push(Edge(picked.t, kv.first, kv.second)); } ret += picked.d; } cout << ret << endl; return 0; }
HackerRank "Prim's (MST) : Special Subtree"
时间: 2024-10-30 00:31:49