4.4.1 Minimum Spanning Tree (Prim)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

/*

For a connected, undirected, weighted graph with possibly negative weights, a minimum spanning tree
(MST) connects all nodes using a subset of edges with minimum possible total weight. On disconnected
input, this implementation instead finds the minimum spanning forest.

Prim's algorithm grows the tree from an arbitrary start node, repeatedly adding the minimum-weight
edge that joins a new node to the current tree, with a priority queue supplying the cheapest such
edge at each step.

- `prim_mst()` populates `mst` with the edges in the minimum spanning forest and returns the total
  MST weight for a global, bidirectionally pre-populated adjacency list `adj`, whose indices
  represent the nodes. Adjacency entries are stored as (`neighbor`, `weight`), while each MST edge
  is stored as (`from`, `to`, `weight`).

The priority queue stores candidate edges as (`weight`, `from`, `to`) and uses `std::greater` to
make it a min-heap. To find a maximum spanning tree instead, use the default max-heap ordering.
Multigraphs are supported; parallel edges are stored as separate adjacency entries and the algorithm
automatically selects the minimum-weight one to each unvisited node.

Time Complexity:
- O(n + m log m) per call, where $n$ is the number of nodes and $m$ is the number of edges.

Space Complexity:
- O(max(n, m)) for storage of the graph, where $n$ is the number of nodes and $m$ is the number of
  edges.
- O(max(n, m)) auxiliary.

*/

#include <cstdint>
#include <functional>
#include <queue>
#include <tuple>
#include <utility>
#include <vector>

std::vector<std::vector<std::pair<int, int>>> adj;  // adj[u] = {(v, weight), ...}
std::vector<std::tuple<int, int, int>> mst;         // (u, v, weight)

int64_t prim_mst() {
  int n = static_cast<int>(adj.size());
  mst.clear();
  std::vector<char> visit(n);
  int64_t total_weight = 0;
  for (int i = 0; i < n; i++) {
    if (visit[i]) {
      continue;
    }
    visit[i] = true;
    using qnode = std::tuple<int, int, int>;  // (weight, u, v)
    std::priority_queue<qnode, std::vector<qnode>, std::greater<>> pq;
    for (auto [v, w] : adj[i]) {
      pq.emplace(w, i, v);
    }
    while (!pq.empty()) {
      auto [w, u, v] = pq.top();
      pq.pop();
      if (visit[v]) {
        continue;
      }
      visit[v] = true;
      mst.emplace_back(u, v, w);
      total_weight += w;  // Overflow warning.
      for (auto [to, ew] : adj[v]) {
        pq.emplace(ew, v, to);
      }
    }
  }
  return total_weight;
}

/*** Example Usage ***/

#include <cassert>
using namespace std;

void add_edge(int u, int v, int w) {
  adj[u].emplace_back(v, w);
  adj[v].emplace_back(u, w);
}

int main() {
  // Two connected components; MST routine returns a minimum spanning forest.
  //      w=4              w=1
  //   0 ----- 1        3 ----- 4
  //    \     /               / |
  // w=3 \   / w=6      w=2 /   | w=4
  //      \ /             / w=3 |
  //       2            5 ----- 6
  adj.assign(7, {});
  add_edge(0, 1, 4);
  add_edge(1, 2, 6);
  add_edge(2, 0, 3);
  add_edge(3, 4, 1);
  add_edge(4, 5, 2);
  add_edge(5, 6, 3);
  add_edge(6, 4, 4);
  assert(prim_mst() == 13);
  assert(
      (mst == vector<tuple<int, int, int>>{{0, 2, 3}, {0, 1, 4}, {3, 4, 1}, {4, 5, 2}, {5, 6, 3}})
  );
  return 0;
}