1use std::cmp::Ordering;
2use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque};
3
4const FLOAT_TOLERANCE: f64 = 1e-9;
5
6type NodeId = usize;
7
8#[derive(Clone, Default)]
9pub struct GraphCore {
10 directed: bool,
11 node_to_id: HashMap<String, NodeId>,
12 id_to_node: Vec<String>,
13 node_order: Vec<NodeId>,
14 succ: Vec<HashSet<NodeId>>,
15 pred: Vec<HashSet<NodeId>>,
16 succ_order: Vec<Vec<NodeId>>,
17 pred_order: Vec<Vec<NodeId>>,
18}
19
20#[derive(Copy, Clone)]
21struct IndexState {
22 cost: f64,
23 node: NodeId,
24}
25
26impl Eq for IndexState {}
27
28impl PartialEq for IndexState {
29 fn eq(&self, other: &Self) -> bool {
30 self.cost.total_cmp(&other.cost) == Ordering::Equal && self.node == other.node
31 }
32}
33
34impl Ord for IndexState {
35 fn cmp(&self, other: &Self) -> Ordering {
36 other
37 .cost
38 .total_cmp(&self.cost)
39 .then_with(|| other.node.cmp(&self.node))
40 }
41}
42
43impl PartialOrd for IndexState {
44 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
45 Some(self.cmp(other))
46 }
47}
48
49impl GraphCore {
50 pub fn new(directed: bool) -> Self {
51 Self {
52 directed,
53 node_to_id: HashMap::new(),
54 id_to_node: Vec::new(),
55 node_order: Vec::new(),
56 succ: Vec::new(),
57 pred: Vec::new(),
58 succ_order: Vec::new(),
59 pred_order: Vec::new(),
60 }
61 }
62
63 pub fn directed(&self) -> bool {
64 self.directed
65 }
66
67 pub fn add_node(&mut self, node: &str) {
68 if self.has_node(node) {
69 return;
70 }
71 let owned = node.to_string();
72 let id = self.id_to_node.len();
73 self.node_to_id.insert(owned.clone(), id);
74 self.id_to_node.push(owned);
75 self.node_order.push(id);
76 self.succ.push(HashSet::new());
77 self.pred.push(HashSet::new());
78 self.succ_order.push(Vec::new());
79 self.pred_order.push(Vec::new());
80 }
81
82 pub fn has_node(&self, node: &str) -> bool {
83 self.node_to_id.contains_key(node)
84 }
85
86 pub fn add_edge(&mut self, left: &str, right: &str) {
87 self.add_node(left);
88 self.add_node(right);
89 let left_id = self
90 .node_to_id
91 .get(left)
92 .copied()
93 .expect("left node should exist");
94 let right_id = self
95 .node_to_id
96 .get(right)
97 .copied()
98 .expect("right node should exist");
99 self.insert_edge_ids(left_id, right_id);
100 }
101
102 pub fn add_nodes(&mut self, nodes: Vec<String>) {
103 for node in nodes {
104 self.add_node(&node);
105 }
106 }
107
108 pub fn add_edges(&mut self, edges: Vec<(String, String)>) {
109 for (left, right) in edges {
110 self.add_edge(&left, &right);
111 }
112 }
113
114 pub fn has_edge(&self, left: &str, right: &str) -> bool {
115 let Some(left_id) = self.node_to_id.get(left).copied() else {
116 return false;
117 };
118 let Some(right_id) = self.node_to_id.get(right).copied() else {
119 return false;
120 };
121 self.succ
122 .get(left_id)
123 .map(|neighbors| neighbors.contains(&right_id))
124 .unwrap_or(false)
125 }
126
127 pub fn remove_edge(&mut self, left: &str, right: &str) {
128 let Some(left_id) = self.node_to_id.get(left).copied() else {
129 return;
130 };
131 let Some(right_id) = self.node_to_id.get(right).copied() else {
132 return;
133 };
134 self.remove_edge_ids(left_id, right_id);
135 }
136
137 pub fn remove_edges(&mut self, edges: Vec<(String, String)>) {
138 for (left, right) in edges {
139 self.remove_edge(&left, &right);
140 }
141 }
142
143 pub fn remove_node(&mut self, node: &str) {
144 let Some(node_id) = self.node_to_id.remove(node) else {
145 return;
146 };
147 let outgoing = self.succ_order[node_id].clone();
148 let incoming = self.pred_order[node_id].clone();
149
150 for neighbor in outgoing {
151 if self.pred[neighbor].remove(&node_id) {
152 remove_ordered_neighbor(&mut self.pred_order[neighbor], node_id);
153 }
154 if !self.directed && self.succ[neighbor].remove(&node_id) {
155 remove_ordered_neighbor(&mut self.succ_order[neighbor], node_id);
156 }
157 }
158
159 for neighbor in incoming {
160 if self.succ[neighbor].remove(&node_id) {
161 remove_ordered_neighbor(&mut self.succ_order[neighbor], node_id);
162 }
163 if !self.directed && self.pred[neighbor].remove(&node_id) {
164 remove_ordered_neighbor(&mut self.pred_order[neighbor], node_id);
165 }
166 }
167
168 self.node_order.retain(|existing| *existing != node_id);
169 self.succ[node_id].clear();
170 self.pred[node_id].clear();
171 self.succ_order[node_id].clear();
172 self.pred_order[node_id].clear();
173 }
174
175 pub fn remove_nodes(&mut self, nodes: Vec<String>) {
176 let removed = nodes
177 .into_iter()
178 .filter_map(|node| self.node_to_id.get(&node).copied())
179 .collect::<HashSet<_>>();
180 if removed.is_empty() {
181 return;
182 }
183 if removed.len() <= 8 || removed.len() * 8 < self.node_order.len() {
184 let names = removed
185 .iter()
186 .map(|id| self.id_to_node[*id].clone())
187 .collect::<Vec<_>>();
188 for name in names {
189 self.remove_node(&name);
190 }
191 return;
192 }
193 for node_id in &removed {
194 self.node_to_id.remove(&self.id_to_node[*node_id]);
195 self.succ[*node_id].clear();
196 self.pred[*node_id].clear();
197 self.succ_order[*node_id].clear();
198 self.pred_order[*node_id].clear();
199 }
200 self.node_order.retain(|node_id| !removed.contains(node_id));
201 for node_id in &self.node_order {
202 self.succ[*node_id].retain(|neighbor| !removed.contains(neighbor));
203 self.pred[*node_id].retain(|neighbor| !removed.contains(neighbor));
204 self.succ_order[*node_id].retain(|neighbor| !removed.contains(neighbor));
205 self.pred_order[*node_id].retain(|neighbor| !removed.contains(neighbor));
206 }
207 }
208
209 pub fn nodes(&self) -> Vec<String> {
210 self.node_order
211 .iter()
212 .map(|node_id| self.id_to_node[*node_id].clone())
213 .collect()
214 }
215
216 pub fn edges(&self) -> Vec<(String, String)> {
217 let mut edges = Vec::with_capacity(
218 self.node_order
219 .iter()
220 .map(|node_id| self.succ_order[*node_id].len())
221 .sum(),
222 );
223 let mut seen = HashSet::new();
224 for left_id in &self.node_order {
225 for right_id in &self.succ_order[*left_id] {
226 if self.directed || seen.insert(canonical_id_edge(*left_id, *right_id)) {
227 edges.push((
228 self.id_to_node[*left_id].clone(),
229 self.id_to_node[*right_id].clone(),
230 ));
231 }
232 }
233 }
234 edges
235 }
236
237 pub fn neighbors(&self, node: &str) -> Vec<String> {
238 self.neighbor_names(self.node_to_id.get(node).copied(), false)
239 }
240
241 pub fn predecessors(&self, node: &str) -> Vec<String> {
242 self.neighbor_names(self.node_to_id.get(node).copied(), true)
243 }
244
245 pub fn successors(&self, node: &str) -> Vec<String> {
246 self.neighbors(node)
247 }
248
249 pub fn all_neighbors(&self, node: &str) -> Vec<String> {
250 let Some(node_id) = self.node_to_id.get(node).copied() else {
251 return Vec::new();
252 };
253 let mut seen = HashSet::new();
254 let mut output = Vec::new();
255 for neighbor in &self.pred_order[node_id] {
256 if seen.insert(*neighbor) {
257 output.push(self.id_to_node[*neighbor].clone());
258 }
259 }
260 for neighbor in &self.succ_order[node_id] {
261 if seen.insert(*neighbor) {
262 output.push(self.id_to_node[*neighbor].clone());
263 }
264 }
265 output
266 }
267
268 pub fn in_degree(&self, node: &str) -> usize {
269 self.node_to_id
270 .get(node)
271 .map(|node_id| self.pred[*node_id].len())
272 .unwrap_or(0)
273 }
274
275 pub fn out_degree(&self, node: &str) -> usize {
276 self.node_to_id
277 .get(node)
278 .map(|node_id| self.succ[*node_id].len())
279 .unwrap_or(0)
280 }
281
282 pub fn clone_graph(&self) -> Self {
283 self.clone()
284 }
285
286 pub fn induced_subgraph(&self, nodes: Vec<String>) -> Self {
287 let keep = nodes
288 .into_iter()
289 .filter_map(|node| self.node_to_id.get(&node).copied())
290 .collect::<HashSet<_>>();
291 let ordered_keep = self
292 .node_order
293 .iter()
294 .copied()
295 .filter(|node_id| keep.contains(node_id))
296 .collect::<Vec<_>>();
297 let mut graph = Self::new(self.directed);
298 let mut id_map = HashMap::new();
299 for old_id in &ordered_keep {
300 let name = &self.id_to_node[*old_id];
301 graph.add_node(name);
302 id_map.insert(
303 *old_id,
304 graph
305 .node_to_id
306 .get(name)
307 .copied()
308 .expect("new node should exist"),
309 );
310 }
311 for old_id in &ordered_keep {
312 let new_id = id_map[old_id];
313 graph.succ[new_id] = self.succ[*old_id]
314 .iter()
315 .filter_map(|neighbor| id_map.get(neighbor).copied())
316 .collect();
317 graph.pred[new_id] = self.pred[*old_id]
318 .iter()
319 .filter_map(|neighbor| id_map.get(neighbor).copied())
320 .collect();
321 graph.succ_order[new_id] = self.succ_order[*old_id]
322 .iter()
323 .filter_map(|neighbor| id_map.get(neighbor).copied())
324 .collect();
325 graph.pred_order[new_id] = self.pred_order[*old_id]
326 .iter()
327 .filter_map(|neighbor| id_map.get(neighbor).copied())
328 .collect();
329 }
330 graph
331 }
332
333 pub fn to_undirected_graph(&self) -> Self {
334 if !self.directed {
335 return self.clone();
336 }
337 let mut graph = Self::new(false);
338 let mut id_map = HashMap::new();
339 for node_id in &self.node_order {
340 let name = &self.id_to_node[*node_id];
341 graph.add_node(name);
342 id_map.insert(
343 *node_id,
344 graph
345 .node_to_id
346 .get(name)
347 .copied()
348 .expect("new node should exist"),
349 );
350 }
351 for left_id in &self.node_order {
352 let new_left = id_map[left_id];
353 for right_id in &self.succ_order[*left_id] {
354 let new_right = id_map[right_id];
355 graph.insert_edge_ids(new_left, new_right);
356 }
357 }
358 graph
359 }
360
361 pub fn connected_components(&self) -> Vec<Vec<String>> {
362 let mut seen = vec![false; self.id_to_node.len()];
363 let mut components = Vec::new();
364
365 for node_id in &self.node_order {
366 if seen[*node_id] {
367 continue;
368 }
369 let mut queue = VecDeque::from([*node_id]);
370 let mut component = Vec::new();
371 seen[*node_id] = true;
372 while let Some(current) = queue.pop_front() {
373 component.push(self.id_to_node[current].clone());
374 for neighbor in &self.pred_order[current] {
375 if !seen[*neighbor] {
376 seen[*neighbor] = true;
377 queue.push_back(*neighbor);
378 }
379 }
380 for neighbor in &self.succ_order[current] {
381 if !seen[*neighbor] {
382 seen[*neighbor] = true;
383 queue.push_back(*neighbor);
384 }
385 }
386 }
387 components.push(component);
388 }
389 components
390 }
391
392 pub fn is_weakly_connected(&self) -> bool {
393 let node_count = self.node_order.len();
394 if node_count <= 1 {
395 return true;
396 }
397 self.connected_components().len() == 1
398 }
399
400 pub fn topological_sort(&self) -> Result<Vec<String>, String> {
401 let mut indegree = vec![0usize; self.id_to_node.len()];
402 for node_id in &self.node_order {
403 indegree[*node_id] = self.pred[*node_id].len();
404 }
405
406 let mut ready = VecDeque::new();
407 for node_id in &self.node_order {
408 if indegree[*node_id] == 0 {
409 ready.push_back(*node_id);
410 }
411 }
412
413 let mut output = Vec::new();
414 while let Some(node_id) = ready.pop_front() {
415 output.push(self.id_to_node[node_id].clone());
416 for neighbor in &self.succ_order[node_id] {
417 indegree[*neighbor] -= 1;
418 if indegree[*neighbor] == 0 {
419 ready.push_back(*neighbor);
420 }
421 }
422 }
423
424 if output.len() != self.node_order.len() {
425 return Err("Graph contains a cycle".to_string());
426 }
427 Ok(output)
428 }
429
430 pub fn strongly_connected_components(&self) -> Vec<Vec<String>> {
433 let unvisited = usize::MAX;
434 let mut index = vec![unvisited; self.id_to_node.len()];
435 let mut low = vec![0usize; self.id_to_node.len()];
436 let mut on_stack = vec![false; self.id_to_node.len()];
437 let mut stack: Vec<NodeId> = Vec::new();
438 let mut next = 0usize;
439 let mut output = Vec::new();
440 for &root in &self.node_order {
441 if index[root] != unvisited {
442 continue;
443 }
444 index[root] = next;
445 low[root] = next;
446 next += 1;
447 stack.push(root);
448 on_stack[root] = true;
449 let mut work: Vec<(NodeId, usize)> = vec![(root, 0)];
450 while let Some(&(node, child)) = work.last() {
451 if let Some(&succ) = self.succ_order[node].get(child) {
452 if let Some(top) = work.last_mut() {
453 top.1 += 1;
454 }
455 if index[succ] == unvisited {
456 index[succ] = next;
457 low[succ] = next;
458 next += 1;
459 stack.push(succ);
460 on_stack[succ] = true;
461 work.push((succ, 0));
462 } else if on_stack[succ] {
463 low[node] = low[node].min(index[succ]);
464 }
465 continue;
466 }
467 work.pop();
468 if let Some(&(parent, _)) = work.last() {
469 low[parent] = low[parent].min(low[node]);
470 }
471 if low[node] == index[node] {
472 let mut component = Vec::new();
473 while let Some(member) = stack.pop() {
474 on_stack[member] = false;
475 component.push(self.id_to_node[member].clone());
476 if member == node {
477 break;
478 }
479 }
480 output.push(component);
481 }
482 }
483 }
484 output
485 }
486
487 pub fn shortest_path(&self, source: &str, target: &str) -> Option<Vec<String>> {
488 let source_id = self.node_to_id.get(source).copied()?;
489 let target_id = self.node_to_id.get(target).copied()?;
490 if source_id == target_id {
491 return Some(vec![source.to_string()]);
492 }
493
494 let mut queue = VecDeque::from([source_id]);
495 let mut visited = vec![false; self.id_to_node.len()];
496 let mut previous = vec![None; self.id_to_node.len()];
497 visited[source_id] = true;
498
499 while let Some(current) = queue.pop_front() {
500 for neighbor in &self.succ_order[current] {
501 if visited[*neighbor] {
502 continue;
503 }
504 visited[*neighbor] = true;
505 previous[*neighbor] = Some(current);
506 if *neighbor == target_id {
507 return Some(self.reconstruct_id_path(&previous, source_id, target_id));
508 }
509 queue.push_back(*neighbor);
510 }
511 }
512
513 None
514 }
515
516 pub fn shortest_path_length(&self, source: &str, target: &str) -> Option<usize> {
517 self.shortest_path(source, target)
518 .map(|path| path.len().saturating_sub(1))
519 }
520
521 pub fn ego_graph_nodes(&self, center: &str, radius: usize) -> Vec<String> {
522 let Some(center_id) = self.node_to_id.get(center).copied() else {
523 return Vec::new();
524 };
525
526 let mut queue = VecDeque::from([(center_id, 0usize)]);
527 let mut visited = vec![false; self.id_to_node.len()];
528 visited[center_id] = true;
529
530 while let Some((current, depth)) = queue.pop_front() {
531 if depth >= radius {
532 continue;
533 }
534 for neighbor in &self.succ_order[current] {
535 if !visited[*neighbor] {
536 visited[*neighbor] = true;
537 queue.push_back((*neighbor, depth + 1));
538 }
539 }
540 }
541
542 self.node_order
543 .iter()
544 .filter(|node_id| visited[**node_id])
545 .map(|node_id| self.id_to_node[*node_id].clone())
546 .collect()
547 }
548
549 pub fn multi_source_dijkstra_path(
550 &self,
551 sources: Vec<String>,
552 weights: Vec<(String, String, f64)>,
553 ) -> Result<Vec<(String, Vec<String>)>, String> {
554 let source_ids = dedupe_preserve_order(
555 sources
556 .into_iter()
557 .map(|source| {
558 self.node_to_id
559 .get(&source)
560 .copied()
561 .ok_or_else(|| format!("Node not found: {source}"))
562 })
563 .collect::<Result<Vec<_>, _>>()?,
564 );
565 if source_ids.is_empty() {
566 return Ok(Vec::new());
567 }
568
569 let weighted_neighbors = self.build_weighted_neighbors(weights)?;
570 let ranks = self.lex_ranks();
571 let mut heap = BinaryHeap::new();
572 let mut distances = vec![f64::INFINITY; self.id_to_node.len()];
573 let mut paths: Vec<Option<Vec<NodeId>>> = vec![None; self.id_to_node.len()];
574
575 for source_id in &source_ids {
576 distances[*source_id] = 0.0;
577 paths[*source_id] = Some(vec![*source_id]);
578 heap.push(IndexState {
579 cost: 0.0,
580 node: *source_id,
581 });
582 }
583
584 while let Some(IndexState { cost, node }) = heap.pop() {
585 if cost > distances[node] + FLOAT_TOLERANCE {
586 continue;
587 }
588 let Some(current_path) = paths[node].clone() else {
589 continue;
590 };
591
592 for (neighbor, weight) in &weighted_neighbors[node] {
593 let next_cost = cost + *weight;
594 let mut next_path = current_path.clone();
595 next_path.push(*neighbor);
596
597 let should_update = if next_cost + FLOAT_TOLERANCE < distances[*neighbor] {
598 true
599 } else if (next_cost - distances[*neighbor]).abs() <= FLOAT_TOLERANCE {
600 match &paths[*neighbor] {
601 None => true,
602 Some(existing_path) => path_less(&next_path, existing_path, &ranks),
603 }
604 } else {
605 false
606 };
607
608 if should_update {
609 distances[*neighbor] = next_cost;
610 paths[*neighbor] = Some(next_path);
611 heap.push(IndexState {
612 cost: next_cost,
613 node: *neighbor,
614 });
615 }
616 }
617 }
618
619 let mut output = Vec::new();
620 for node_id in &self.node_order {
621 let Some(path) = &paths[*node_id] else {
622 continue;
623 };
624 output.push((
625 self.id_to_node[*node_id].clone(),
626 path.iter()
627 .map(|path_id| self.id_to_node[*path_id].clone())
628 .collect(),
629 ));
630 }
631 Ok(output)
632 }
633
634 pub fn steiner_tree_nodes(
635 &self,
636 terminals: Vec<String>,
637 weights: Vec<(String, String, f64)>,
638 ) -> Result<Vec<String>, String> {
639 let terminal_ids = dedupe_preserve_order(
640 terminals
641 .into_iter()
642 .map(|terminal| {
643 self.node_to_id
644 .get(&terminal)
645 .copied()
646 .ok_or_else(|| format!("Node not found: {terminal}"))
647 })
648 .collect::<Result<Vec<_>, _>>()?,
649 );
650 if terminal_ids.is_empty() {
651 return Ok(Vec::new());
652 }
653 if terminal_ids.len() == 1 {
654 return Ok(vec![self.id_to_node[terminal_ids[0]].clone()]);
655 }
656
657 let weighted_neighbors = self.build_weighted_neighbors(weights)?;
658 let ranks = self.lex_ranks();
659 let mut metric_edges: Vec<(f64, NodeId, NodeId, Vec<NodeId>)> = Vec::new();
660
661 for index in 0..terminal_ids.len() {
662 let left = terminal_ids[index];
663 let targets = &terminal_ids[index + 1..];
664 let shortest_paths = self.weighted_shortest_paths_to_target_indices(
665 left,
666 targets,
667 &weighted_neighbors,
668 &ranks,
669 );
670 for right in targets {
671 let Some((distance, path)) = shortest_paths.get(right) else {
672 return Err(format!(
673 "No path between {} and {}",
674 self.id_to_node[left], self.id_to_node[*right]
675 ));
676 };
677 metric_edges.push((*distance, left, *right, path.clone()));
678 }
679 }
680
681 metric_edges.sort_by(|left, right| {
682 left.0
683 .total_cmp(&right.0)
684 .then_with(|| ranks[left.1].cmp(&ranks[right.1]))
685 .then_with(|| ranks[left.2].cmp(&ranks[right.2]))
686 });
687
688 let mut metric_dsu = IndexDisjointSet::new(self.id_to_node.len());
689 let mut expanded_nodes = terminal_ids.iter().copied().collect::<HashSet<_>>();
690 for (_distance, left, right, path) in metric_edges {
691 if metric_dsu.union(left, right) {
692 expanded_nodes.extend(path);
693 }
694 }
695
696 Ok(self.finalize_steiner_tree(
697 terminal_ids,
698 expanded_nodes,
699 &weighted_neighbors,
700 ))
701 }
702
703 fn insert_edge_ids(&mut self, left_id: NodeId, right_id: NodeId) {
704 if self.succ[left_id].insert(right_id) {
705 self.succ_order[left_id].push(right_id);
706 }
707 if self.pred[right_id].insert(left_id) {
708 self.pred_order[right_id].push(left_id);
709 }
710 if !self.directed {
711 if self.succ[right_id].insert(left_id) {
712 self.succ_order[right_id].push(left_id);
713 }
714 if self.pred[left_id].insert(right_id) {
715 self.pred_order[left_id].push(right_id);
716 }
717 }
718 }
719
720 fn remove_edge_ids(&mut self, left_id: NodeId, right_id: NodeId) {
721 if self.succ[left_id].remove(&right_id) {
722 remove_ordered_neighbor(&mut self.succ_order[left_id], right_id);
723 }
724 if self.pred[right_id].remove(&left_id) {
725 remove_ordered_neighbor(&mut self.pred_order[right_id], left_id);
726 }
727 if !self.directed {
728 if self.succ[right_id].remove(&left_id) {
729 remove_ordered_neighbor(&mut self.succ_order[right_id], left_id);
730 }
731 if self.pred[left_id].remove(&right_id) {
732 remove_ordered_neighbor(&mut self.pred_order[left_id], right_id);
733 }
734 }
735 }
736
737 fn neighbor_names(&self, node_id: Option<NodeId>, reverse: bool) -> Vec<String> {
738 let Some(node_id) = node_id else {
739 return Vec::new();
740 };
741 let neighbors = if reverse {
742 &self.pred_order[node_id]
743 } else {
744 &self.succ_order[node_id]
745 };
746 neighbors
747 .iter()
748 .map(|neighbor| self.id_to_node[*neighbor].clone())
749 .collect()
750 }
751
752 fn reconstruct_id_path(
753 &self,
754 previous: &[Option<NodeId>],
755 source: NodeId,
756 target: NodeId,
757 ) -> Vec<String> {
758 let mut path = vec![self.id_to_node[target].clone()];
759 let mut current = target;
760 while current != source {
761 let Some(parent) = previous[current] else {
762 break;
763 };
764 current = parent;
765 path.push(self.id_to_node[current].clone());
766 }
767 path.reverse();
768 path
769 }
770
771 fn lex_ranks(&self) -> Vec<usize> {
772 let mut ranked = self.node_order.clone();
773 ranked.sort_by(|left, right| self.id_to_node[*left].cmp(&self.id_to_node[*right]));
774 let mut ranks = vec![usize::MAX; self.id_to_node.len()];
775 for (rank, node_id) in ranked.into_iter().enumerate() {
776 ranks[node_id] = rank;
777 }
778 ranks
779 }
780
781 fn build_weighted_neighbors(
782 &self,
783 weights: Vec<(String, String, f64)>,
784 ) -> Result<Vec<Vec<(NodeId, f64)>>, String> {
785 let mut overrides =
786 HashMap::with_capacity(weights.len() * if self.directed { 1 } else { 2 });
787 for (left, right, weight) in weights {
788 let Some(left_id) = self.node_to_id.get(&left).copied() else {
789 return Err(format!("Node not found: {left}"));
790 };
791 let Some(right_id) = self.node_to_id.get(&right).copied() else {
792 return Err(format!("Node not found: {right}"));
793 };
794 overrides.insert((left_id, right_id), weight);
795 if !self.directed {
796 overrides.insert((right_id, left_id), weight);
797 }
798 }
799
800 let mut weighted_neighbors = vec![Vec::<(NodeId, f64)>::new(); self.id_to_node.len()];
801 for node_id in &self.node_order {
802 let mut neighbors = Vec::with_capacity(self.succ_order[*node_id].len());
803 for neighbor in &self.succ_order[*node_id] {
804 neighbors.push((
805 *neighbor,
806 overrides
807 .get(&(*node_id, *neighbor))
808 .copied()
809 .unwrap_or(1.0),
810 ));
811 }
812 weighted_neighbors[*node_id] = neighbors;
813 }
814 Ok(weighted_neighbors)
815 }
816
817 fn finalize_steiner_tree(
818 &self,
819 terminals: Vec<NodeId>,
820 expanded_nodes: HashSet<NodeId>,
821 weighted_neighbors: &[Vec<(NodeId, f64)>],
822 ) -> Vec<String> {
823 let ordered_expanded_nodes = self
824 .node_order
825 .iter()
826 .copied()
827 .filter(|node_id| expanded_nodes.contains(node_id))
828 .collect::<Vec<_>>();
829 let mut position = vec![usize::MAX; self.id_to_node.len()];
830 for (index, node_id) in ordered_expanded_nodes.iter().enumerate() {
831 position[*node_id] = index;
832 }
833
834 let mut original_edges = Vec::new();
835 let mut seen_edges = HashSet::new();
836 for left_id in &ordered_expanded_nodes {
837 for (right_id, weight) in &weighted_neighbors[*left_id] {
838 if position[*right_id] == usize::MAX {
839 continue;
840 }
841 if self.directed || seen_edges.insert(canonical_id_edge(*left_id, *right_id)) {
842 original_edges.push((*weight, *left_id, *right_id));
843 }
844 }
845 }
846 let ranks = self.lex_ranks();
847 original_edges.sort_by(|left, right| {
848 left.0
849 .total_cmp(&right.0)
850 .then_with(|| ranks[left.1].cmp(&ranks[right.1]))
851 .then_with(|| ranks[left.2].cmp(&ranks[right.2]))
852 });
853
854 let mut tree = vec![HashSet::<NodeId>::new(); self.id_to_node.len()];
855 let mut induced_dsu = IndexDisjointSet::new(self.id_to_node.len());
856 for (_weight, left, right) in original_edges {
857 if induced_dsu.union(left, right) {
858 tree[left].insert(right);
859 tree[right].insert(left);
860 }
861 }
862
863 let terminals_set = terminals.into_iter().collect::<HashSet<_>>();
864 let mut removable = ordered_expanded_nodes
865 .iter()
866 .filter_map(|node_id| {
867 if !terminals_set.contains(node_id) && tree[*node_id].len() <= 1 {
868 Some(*node_id)
869 } else {
870 None
871 }
872 })
873 .collect::<VecDeque<_>>();
874 let mut removed = vec![false; self.id_to_node.len()];
875 while let Some(node_id) = removable.pop_front() {
876 if removed[node_id] || terminals_set.contains(&node_id) || tree[node_id].len() > 1 {
877 continue;
878 }
879 removed[node_id] = true;
880 let neighbors = std::mem::take(&mut tree[node_id]);
881 for neighbor in neighbors {
882 if tree[neighbor].remove(&node_id)
883 && !terminals_set.contains(&neighbor)
884 && tree[neighbor].len() <= 1
885 {
886 removable.push_back(neighbor);
887 }
888 }
889 }
890
891 self.node_order
892 .iter()
893 .filter(|node_id| {
894 expanded_nodes.contains(node_id)
895 && !removed[**node_id]
896 && (!tree[**node_id].is_empty() || terminals_set.contains(node_id))
897 })
898 .map(|node_id| self.id_to_node[*node_id].clone())
899 .collect()
900 }
901
902 fn weighted_shortest_paths_to_target_indices(
903 &self,
904 source: NodeId,
905 targets: &[NodeId],
906 weighted_neighbors: &[Vec<(NodeId, f64)>],
907 ranks: &[usize],
908 ) -> HashMap<NodeId, (f64, Vec<NodeId>)> {
909 if targets.is_empty() {
910 return HashMap::new();
911 }
912 let mut remaining = targets.iter().copied().collect::<HashSet<_>>();
913 let mut heap = BinaryHeap::new();
914 let mut distances = vec![f64::INFINITY; self.id_to_node.len()];
915 let mut paths: Vec<Option<Vec<NodeId>>> = vec![None; self.id_to_node.len()];
916
917 distances[source] = 0.0;
918 paths[source] = Some(vec![source]);
919 heap.push(IndexState {
920 cost: 0.0,
921 node: source,
922 });
923
924 while let Some(IndexState { cost, node }) = heap.pop() {
925 if cost > distances[node] + FLOAT_TOLERANCE {
926 continue;
927 }
928 remaining.remove(&node);
929 if remaining.is_empty() {
930 break;
931 }
932 let Some(current_path) = paths[node].clone() else {
933 continue;
934 };
935 for (neighbor, weight) in &weighted_neighbors[node] {
936 let next_cost = cost + *weight;
937 let mut next_path = current_path.clone();
938 next_path.push(*neighbor);
939 let should_update = if next_cost + FLOAT_TOLERANCE < distances[*neighbor] {
940 true
941 } else if (next_cost - distances[*neighbor]).abs() <= FLOAT_TOLERANCE {
942 match &paths[*neighbor] {
943 None => true,
944 Some(existing_path) => path_less(&next_path, existing_path, ranks),
945 }
946 } else {
947 false
948 };
949 if should_update {
950 distances[*neighbor] = next_cost;
951 paths[*neighbor] = Some(next_path);
952 heap.push(IndexState {
953 cost: next_cost,
954 node: *neighbor,
955 });
956 }
957 }
958 }
959
960 targets
961 .iter()
962 .filter_map(|target| {
963 let path = paths[*target].clone()?;
964 Some((*target, (distances[*target], path)))
965 })
966 .collect()
967 }
968}
969
970struct IndexDisjointSet {
971 parent: Vec<NodeId>,
972 rank: Vec<usize>,
973}
974
975impl IndexDisjointSet {
976 fn new(size: usize) -> Self {
977 Self {
978 parent: (0..size).collect(),
979 rank: vec![0; size],
980 }
981 }
982
983 fn find(&mut self, node: NodeId) -> NodeId {
984 if self.parent[node] != node {
985 let root = self.find(self.parent[node]);
986 self.parent[node] = root;
987 }
988 self.parent[node]
989 }
990
991 fn union(&mut self, left: NodeId, right: NodeId) -> bool {
992 let left_root = self.find(left);
993 let right_root = self.find(right);
994 if left_root == right_root {
995 return false;
996 }
997 let left_rank = self.rank[left_root];
998 let right_rank = self.rank[right_root];
999 if left_rank < right_rank {
1000 self.parent[left_root] = right_root;
1001 } else if left_rank > right_rank {
1002 self.parent[right_root] = left_root;
1003 } else {
1004 self.parent[right_root] = left_root;
1005 self.rank[left_root] += 1;
1006 }
1007 true
1008 }
1009}
1010
1011fn dedupe_preserve_order<T: Eq + std::hash::Hash + Copy>(values: Vec<T>) -> Vec<T> {
1012 let mut seen = HashSet::new();
1013 let mut output = Vec::new();
1014 for value in values {
1015 if seen.insert(value) {
1016 output.push(value);
1017 }
1018 }
1019 output
1020}
1021
1022fn canonical_id_edge(left: NodeId, right: NodeId) -> (NodeId, NodeId) {
1023 if left <= right {
1024 (left, right)
1025 } else {
1026 (right, left)
1027 }
1028}
1029
1030fn remove_ordered_neighbor(order: &mut Vec<NodeId>, neighbor: NodeId) {
1031 order.retain(|entry| *entry != neighbor);
1032}
1033
1034fn path_less(left: &[NodeId], right: &[NodeId], ranks: &[usize]) -> bool {
1035 for (left_node, right_node) in left.iter().zip(right.iter()) {
1036 match ranks[*left_node].cmp(&ranks[*right_node]) {
1037 Ordering::Less => return true,
1038 Ordering::Greater => return false,
1039 Ordering::Equal => {}
1040 }
1041 }
1042 left.len() < right.len()
1043}