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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
//! Maximum-weight spanning **branching** (rooted forest) over a directed graph,
//! via Chu–Liu/Edmonds.
//!
//! A branching is a set of arcs with **in-degree ≤ 1** at every node and **no
//! cycles** — equivalently, a forest of arborescences, each edge pointing away
//! from its tree's root. [`max_branching`] returns the branching of maximum total
//! arc weight, subject to a per-node **`root_affinity`**: the value a node gets for
//! staying a root (taking no parent). A node adopts a real parent only when its best
//! incoming arc beats its `root_affinity`, so `root_affinity` is the knob that trades
//! one big tree (low affinity) against many small ones (high affinity).
//!
//! Implemented by adding a virtual super-root with an arc of weight `root_affinity[v]`
//! into every node, then computing a maximum spanning **arborescence** from it (every
//! node reachable, in-degree 1) — the super-root's children are the forest roots.
//! At `K ≤ few hundred` nodes the cubic-ish contraction cost is negligible.
use std::collections::HashSet;
/// A maximum-weight rooted forest over `n` nodes.
pub struct Branching {
/// `parent[v] = Some(u)` for the chosen arc `u → v`; `None` when `v` is a root.
pub parent: Vec<Option<usize>>,
/// Component (tree) id per node, dense in `0..roots.len()`, ordered by ascending root.
pub tree: Vec<usize>,
/// Root node ids (in-degree 0), ascending.
pub roots: Vec<usize>,
}
/// One directed arc in the contracted working graph. `orig` indexes the original
/// (super-root-augmented) arc list and is stable across contraction levels; `landed`
/// is the arc's target in the *current* level's numbering (which differs from `v` only
/// for an arc collapsed onto a cycle super-node — there `landed` is the pre-collapse
/// cycle member, needed to know where to break the cycle on expansion).
struct Arc {
u: usize,
v: usize,
w: f64,
orig: usize,
landed: usize,
}
/// Maximum-weight spanning branching. `arcs` are `(u, v, weight)` directed `u → v`
/// over nodes `0..n`; `root_affinity[v]` is the score for leaving `v` a root. Arcs into
/// a node compete with that node's `root_affinity`. Ties resolve to the earliest arc.
pub fn max_branching(n: usize, arcs: &[(usize, usize, f32)], root_affinity: &[f32]) -> Branching {
assert_eq!(root_affinity.len(), n, "root_affinity must have length n");
if n == 0 {
return Branching {
parent: Vec::new(),
tree: Vec::new(),
roots: Vec::new(),
};
}
// Super-root = node `n`, with an arc n → v of weight root_affinity[v] into every v.
let sroot = n;
let mut all: Vec<Arc> = Vec::with_capacity(arcs.len() + n);
for (i, &(u, v, w)) in arcs.iter().enumerate() {
if u == v {
continue; // self-loops never help
}
all.push(Arc {
u,
v,
w: w as f64,
orig: i,
landed: v,
});
}
let sroot_orig_base = arcs.len();
for (v, &aff) in root_affinity.iter().enumerate() {
all.push(Arc {
u: sroot,
v,
w: aff as f64,
orig: sroot_orig_base + v,
landed: v,
});
}
let used = solve(n + 1, sroot, &all);
// Reconstruct parent pointers from the chosen original arc ids.
let mut parent: Vec<Option<usize>> = vec![None; n];
for &orig in &used {
let (u, v) = if orig < sroot_orig_base {
(arcs[orig].0, arcs[orig].1)
} else {
(sroot, orig - sroot_orig_base)
};
if v < n {
parent[v] = if u == sroot { None } else { Some(u) };
}
}
// Roots = in-degree-0 nodes; component ids by ascending root.
let roots: Vec<usize> = (0..n).filter(|&v| parent[v].is_none()).collect();
let mut root_of_comp = vec![usize::MAX; n];
for (c, &r) in roots.iter().enumerate() {
root_of_comp[r] = c;
}
let mut tree = vec![usize::MAX; n];
for v in 0..n {
// walk to the root, then paint the path with its component id
let mut path = Vec::new();
let mut x = v;
while tree[x] == usize::MAX && root_of_comp[x] == usize::MAX {
path.push(x);
match parent[x] {
Some(p) => x = p,
None => break,
}
}
let comp = if tree[x] != usize::MAX {
tree[x]
} else {
root_of_comp[x]
};
for &p in &path {
tree[p] = comp;
}
tree[v] = comp;
}
Branching {
parent,
tree,
roots,
}
}
/// Maximum spanning arborescence rooted at `root` over `n` nodes; returns the chosen
/// original arc ids. Every non-root node is assumed reachable (the super-root guarantees
/// this at the top level, and every contracted cycle keeps its super-root arc).
fn solve(n: usize, root: usize, arcs: &[Arc]) -> Vec<usize> {
// Best (max-weight) incoming arc per node.
let mut best: Vec<Option<usize>> = vec![None; n];
for (i, a) in arcs.iter().enumerate() {
if a.v == root {
continue;
}
match best[a.v] {
None => best[a.v] = Some(i),
Some(j) if a.w > arcs[j].w => best[a.v] = Some(i),
_ => {}
}
}
// Find a cycle in the functional parent graph (v → arcs[best[v]].u).
let par = |v: usize| best[v].map(|i| arcs[i].u);
let mut color = vec![0u8; n]; // 0 white, 1 gray, 2 black
let mut cycle: Option<Vec<usize>> = None;
for s in 0..n {
if color[s] != 0 || s == root {
continue;
}
let mut stack: Vec<usize> = Vec::new();
let mut v = s;
loop {
if v == root || color[v] == 2 {
break;
}
if color[v] == 1 {
let start = stack.iter().position(|&x| x == v).unwrap();
cycle = Some(stack[start..].to_vec());
break;
}
color[v] = 1;
stack.push(v);
match par(v) {
Some(p) => v = p,
None => break,
}
}
for &x in &stack {
color[x] = 2;
}
if cycle.is_some() {
break;
}
}
// No cycle: the best arcs already form the arborescence.
let Some(cyc) = cycle else {
let mut used = Vec::new();
for (v, slot) in best.iter().enumerate() {
if v == root {
continue;
}
if let Some(i) = slot {
used.push(arcs[*i].orig);
}
}
return used;
};
// Contract the cycle into a single super-node `cnode`.
let in_cycle: HashSet<usize> = cyc.iter().copied().collect();
let mut map = vec![usize::MAX; n];
let mut next = 0;
for (v, m) in map.iter_mut().enumerate() {
if !in_cycle.contains(&v) {
*m = next;
next += 1;
}
}
let cnode = next;
next += 1;
for &v in &cyc {
map[v] = cnode;
}
let new_n = next;
let mut new_arcs: Vec<Arc> = Vec::with_capacity(arcs.len());
for a in arcs {
let nu = map[a.u];
let nv = map[a.v];
if nu == nv {
continue; // internal to the cycle (or self) — drop
}
let w = if in_cycle.contains(&a.v) {
// reduced weight for maximum: gain of this arc minus the cycle arc it replaces
a.w - arcs[best[a.v].unwrap()].w
} else {
a.w
};
new_arcs.push(Arc {
u: nu,
v: nv,
w,
orig: a.orig,
landed: a.v,
});
}
let sub_used = solve(new_n, map[root], &new_arcs);
let sub_set: HashSet<usize> = sub_used.iter().copied().collect();
// The one used arc entering `cnode` breaks the cycle at `v_enter`.
let mut v_enter = None;
for a in &new_arcs {
if a.v == cnode && sub_set.contains(&a.orig) {
v_enter = Some(a.landed);
break;
}
}
let v_enter = v_enter.expect("contracted cycle must have an entering arc");
let mut used = sub_used;
for &x in &cyc {
if x != v_enter {
used.push(arcs[best[x].unwrap()].orig);
}
}
used
}
#[cfg(test)]
mod tests;