4.4.4 Minimum Spanning Arborescence (Edmonds)

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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

/*

Given a weighted, directed graph and a root node, a spanning arborescence is a set of edges forming
a tree in which every other node is reachable from the root along directed edges. Equivalently,
every node except the root has exactly one incoming edge. The minimum spanning arborescence
minimizes the total edge weight and is the directed analog of the minimum spanning tree.

Edmonds' (a.k.a. the Chu-Liu/Edmonds) algorithm computes the arborescence by first selecting, for
each non-root node, its cheapest incoming edge. If these choices form no cycle, they already
constitute the answer. Otherwise each cycle is contracted into a single supernode, and every edge
entering the cycle has its weight reduced by the cost of the edge it would replace inside the cycle.
Solving the problem on the contracted graph and expanding the cycles back yields the optimum. The
implementation below repeats this contraction in rounds, accumulating the selected weights, until no
cycle remains.

- `directed_mst(n, edges, root)` returns the total weight of the minimum spanning arborescence
  rooted at `root` over `n` nodes numbered $[0, `n`)$, or `std::nullopt` if no arborescence exists.
  Edges are given as (`from`, `to`, `weight`) triples; parallel edges and self-loops are allowed,
  with self-loops simply ignored.

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

Space Complexity:
- O(n + m) auxiliary.

*/

#include <cstdint>
#include <optional>
#include <tuple>
#include <vector>

using Edge = std::tuple<int, int, int64_t>;  // (u, v, weight)

std::optional<int64_t> directed_mst(int n, std::vector<Edge> edges, int root) {
  const int64_t INF = INT64_MAX / 4;
  int64_t total_dist = 0;
  while (true) {
    std::vector<int64_t> min_in(n, INF);
    std::vector<int> pre(n, -1);
    for (const auto &[u, v, w] : edges) {
      if (u != v && w < min_in[v]) {
        min_in[v] = w;
        pre[v] = u;
      }
    }
    for (int v = 0; v < n; v++) {
      if (v != root && min_in[v] == INF) {
        return std::nullopt;
      }
    }
    // Identify cycles formed by the chosen incoming edges, labeling each with an ID in comp.
    int cycles = 0;
    std::vector<int> comp(n, -1), seen(n, -1);
    min_in[root] = 0;
    for (int v = 0; v < n; v++) {
      total_dist += min_in[v];  // Overflow warning.
      int u = v;
      while (seen[u] != v && comp[u] == -1 && u != root) {
        seen[u] = v;
        u = pre[u];
      }
      if (u != root && comp[u] == -1) {
        for (int w = pre[u]; w != u; w = pre[w]) {
          comp[w] = cycles;
        }
        comp[u] = cycles++;
      }
    }
    if (cycles == 0) {
      break;  // No cycle remains; the selected edges form the arborescence.
    }
    // Contract each cycle into a supernode and reduce the weights of edges entering it.
    for (int v = 0; v < n; v++) {
      if (comp[v] == -1) {
        comp[v] = cycles++;
      }
    }
    for (auto &[u, v, w] : edges) {
      int cu = comp[u], cv = comp[v];
      if (cu != cv) {
        w -= min_in[v];  // Overflow warning.
      }
      u = cu;
      v = cv;
    }
    n = cycles;
    root = comp[root];
  }
  return total_dist;
}

/*** Example Usage ***/

#include <cassert>
using namespace std;

int main() {
  //          w=1
  // root=0 -----> 1 <-----+
  //      |      / |       |
  //      |   w=2  |       |
  //   w=5|  /     |w=2    |w=4
  //      v v      v       |
  //      2 -----> 3 ------+
  //          w=3
  vector<Edge> edges{
      {0, 1, 1}, {0, 2, 5}, {1, 2, 2}, {2, 3, 3}, {3, 1, 4},
  };
  assert(directed_mst(4, edges, 0) == 6);  // Cheapest edges 0->1->2->3 form a tree of weight 6.

  //        w=1    w=2
  // root=0 ---> 1 ---> 2       3
  vector<Edge> disconnected{{0, 1, 1}, {1, 2, 1}};
  assert(!directed_mst(4, disconnected, 0));  // Node 3 is unreachable from root 0.

  vector<Edge> negative{{0, 1, -1}};
  assert(directed_mst(2, negative, 0) == -1);  // A weight of -1 is a valid result.
  return 0;
}