Skip to main content

trilogy_parser/
graph.rs

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    /// Tarjan's strongly connected components, sinks first: each component
431    /// follows every component it reaches.
432    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}