4.4.2 Minimum Spanning Tree (Kruskal)

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

/*

Given a connected, undirected, weighted graph with possibly negative weights, its minimum spanning
tree (MST) is a subgraph which is a tree that connects all nodes with a subset of its edges such
that their total weight is minimized. If the input graph is not connected, then this implementation
will find the minimum spanning forest.

Kruskal's algorithm scans the edges in nondecreasing weight order, adding each edge whose endpoints
lie in different components (tracked with a disjoint-set structure) and skipping any edge that would
form a cycle.

- `kruskal_mst(n)` populates `mst` with the edge IDs in the minimum spanning tree (returning the
  total MST weight) for a global, pre-populated edge list `edges` whose endpoints must be numbered
  [0, `n`). Each edge is stored as (`weight`, `u`, `v`).

Time Complexity:
- O(m log n) per call to `kruskal_mst()`, where $m$ is the number of edges and $n$ is the number of
  nodes. The internal DSU uses path compression and union by size, making each DSU operation
  O(alpha(n)) amortized.

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(n) auxiliary stack space for `kruskal_mst()`.

*/

#include <algorithm>
#include <numeric>
#include <tuple>
#include <utility>
#include <vector>

std::vector<std::tuple<int, int, int>> edges;  // (weight, u, v)
std::vector<int> dsu_root, dsu_size, mst;

int find_root(int u) {
  if (dsu_root[u] != u) {
    dsu_root[u] = find_root(dsu_root[u]);
  }
  return dsu_root[u];
}

int kruskal_mst(int n) {
  mst.clear();
  std::vector<int> order(edges.size());
  std::iota(order.begin(), order.end(), 0);
  std::sort(order.begin(), order.end(), [](int a, int b) {
    return std::get<0>(edges[a]) < std::get<0>(edges[b]);
  });
  int total_dist = 0;
  dsu_root.assign(n, 0);
  dsu_size.assign(n, 1);
  std::iota(dsu_root.begin(), dsu_root.end(), 0);
  for (int id : order) {
    auto &[w, a, b] = edges[id];
    int u = find_root(a), v = find_root(b);
    if (u != v) {
      if (dsu_size[u] < dsu_size[v]) {
        std::swap(u, v);
      }
      dsu_root[v] = u;
      dsu_size[u] += dsu_size[v];
      mst.push_back(id);
      total_dist += w;
    }
  }
  return total_dist;
}

/*** Example Usage and Output:

Total distance: 13
3 <-> 4
4 <-> 5
2 <-> 0
5 <-> 6
0 <-> 1

***/

#include <iostream>
using namespace std;

void add_edge(int u, int v, int w) {
  edges.emplace_back(w, u, v);
}

int main() {
  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);
  cout << "Total distance: " << kruskal_mst(7) << endl;
  for (int id : mst) {
    auto &[w, u, v] = edges[id];
    cout << u << " <-> " << v << endl;
  }
  return 0;
}