toolbox_rs/
tarjan.rs

1use log::info;
2
3use crate::graph::{Graph, NodeID};
4use core::cmp::min;
5
6#[derive(Clone)]
7struct DFSNode {
8    caller: NodeID,
9    index: usize,
10    lowlink: NodeID,
11    neighbor: usize,
12    on_stack: bool,
13}
14
15impl DFSNode {
16    pub fn new() -> Self {
17        DFSNode {
18            caller: NodeID::MAX,
19            index: usize::MAX,
20            lowlink: NodeID::MAX,
21            neighbor: usize::MAX,
22            on_stack: false,
23        }
24    }
25}
26
27#[derive(Default)]
28pub struct Tarjan {
29    dfs_state: Vec<DFSNode>,
30    tarjan_stack: Vec<NodeID>,
31}
32
33impl Tarjan {
34    pub fn new() -> Self {
35        Self {
36            dfs_state: Vec::new(),
37            tarjan_stack: Vec::new(),
38        }
39    }
40
41    // TODO: consider adding handlers for small/large SCCs
42    pub fn run<T>(&mut self, graph: &(impl Graph<T> + 'static)) -> Vec<usize> {
43        let mut assignment = Vec::new();
44        assignment.resize(graph.number_of_nodes(), usize::MAX);
45        let mut index = 0;
46        let mut num_scc = 0;
47
48        self.dfs_state
49            .resize(graph.number_of_nodes(), DFSNode::new());
50
51        // assign each node to an SCC if not yet done
52        for n in graph.node_range() {
53            if self.dfs_state[n].index != usize::MAX {
54                continue;
55            }
56            // TODO: consider moving the following to a function to save indentation
57
58            self.stack_push(n, usize::MAX, index);
59            index += 1;
60            let mut last = n;
61
62            loop {
63                if self.dfs_state[last].neighbor < graph.out_degree(last) {
64                    let e = graph
65                        .edge_range(last)
66                        .nth(self.dfs_state[last].neighbor)
67                        .expect("edge range exhausted");
68                    let w = graph.target(e);
69                    self.dfs_state[last].neighbor += 1;
70                    if self.dfs_state[w].index == usize::MAX {
71                        self.stack_push(w, last, index);
72                        index += 1;
73                        last = w;
74                    } else if self.dfs_state[w].on_stack {
75                        self.dfs_state[last].lowlink =
76                            min(self.dfs_state[last].lowlink, self.dfs_state[w].index);
77                    }
78                } else {
79                    if self.dfs_state[last].lowlink == self.dfs_state[last].index {
80                        num_scc += 1;
81                        let mut size = 0;
82                        loop {
83                            let top = self.tarjan_stack.pop().expect("tarjan_stack empty");
84                            self.dfs_state[top].on_stack = false;
85                            size += 1;
86                            assignment[top] = num_scc;
87                            if top == last {
88                                break;
89                            }
90                        }
91                        // TODO: call handler for small/large SCCs
92                        info!("detected SCC of size {size}");
93                    }
94
95                    let new_last = self.dfs_state[last].caller;
96                    if new_last != usize::MAX {
97                        self.dfs_state[new_last].lowlink = min(
98                            self.dfs_state[new_last].lowlink,
99                            self.dfs_state[last].lowlink,
100                        );
101                        last = new_last;
102                    } else {
103                        debug_assert!(n == last);
104                        break;
105                    }
106                }
107            }
108        }
109        assignment
110    }
111
112    fn stack_push(&mut self, w: usize, caller: usize, index: usize) {
113        self.dfs_state[w].caller = caller;
114        self.dfs_state[w].neighbor = 0;
115        self.dfs_state[w].index = index;
116        self.dfs_state[w].lowlink = index;
117        self.dfs_state[w].on_stack = true;
118        self.tarjan_stack.push(w);
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use crate::edge::InputEdge;
125    use crate::static_graph::StaticGraph;
126    use crate::tarjan::Tarjan;
127
128    #[test]
129    fn scc_wiki1() {
130        type Graph = StaticGraph<i32>;
131        let edges = vec![
132            InputEdge::new(0, 1, 3),
133            InputEdge::new(1, 2, 3),
134            InputEdge::new(1, 4, 1),
135            InputEdge::new(1, 5, 6),
136            InputEdge::new(2, 3, 2),
137            InputEdge::new(2, 6, 2),
138            InputEdge::new(3, 2, 2),
139            InputEdge::new(3, 7, 2),
140            InputEdge::new(4, 0, 2),
141            InputEdge::new(4, 5, 2),
142            InputEdge::new(5, 6, 2),
143            InputEdge::new(6, 5, 2),
144            InputEdge::new(7, 3, 2),
145            InputEdge::new(7, 6, 2),
146        ];
147        let graph = Graph::new(edges);
148
149        let mut tarjan = Tarjan::new();
150        assert_eq!(vec![3, 3, 2, 2, 3, 1, 1, 2], tarjan.run(&graph));
151    }
152
153    #[test]
154    fn geekforgeeks() {
155        type Graph = StaticGraph<i32>;
156        let edges = vec![
157            InputEdge::new(1, 0, 3),
158            InputEdge::new(0, 3, 3),
159            InputEdge::new(0, 2, 1),
160            InputEdge::new(2, 1, 6),
161            InputEdge::new(3, 4, 2),
162        ];
163        let graph = Graph::new(edges);
164
165        let mut tarjan = Tarjan::new();
166        assert_eq!(vec![3, 3, 3, 2, 1], tarjan.run(&graph));
167    }
168
169    #[test]
170    fn stanford2() {
171        type Graph = StaticGraph<i32>;
172        let edges = vec![
173            InputEdge::new(0, 6, 3),
174            InputEdge::new(6, 3, 3),
175            InputEdge::new(3, 0, 1),
176            InputEdge::new(6, 8, 6),
177            InputEdge::new(8, 5, 2),
178            InputEdge::new(5, 2, 2),
179            InputEdge::new(2, 8, 2),
180            InputEdge::new(5, 7, 2),
181            InputEdge::new(7, 1, 2),
182            InputEdge::new(4, 7, 2),
183            InputEdge::new(1, 4, 2),
184        ];
185        let graph = Graph::new(edges);
186
187        let mut tarjan = Tarjan::new();
188        assert_eq!(vec![3, 1, 2, 3, 1, 2, 3, 1, 2], tarjan.run(&graph));
189    }
190
191    #[test]
192    fn web1() {
193        type Graph = StaticGraph<i32>;
194        let edges = vec![
195            InputEdge::new(0, 1, 3),
196            InputEdge::new(1, 3, 3),
197            InputEdge::new(1, 4, 1),
198            InputEdge::new(1, 2, 6),
199            InputEdge::new(2, 5, 2),
200            InputEdge::new(4, 1, 2),
201            InputEdge::new(4, 5, 2),
202            InputEdge::new(4, 6, 2),
203            InputEdge::new(5, 7, 2),
204            InputEdge::new(6, 7, 2),
205            InputEdge::new(6, 8, 2),
206            InputEdge::new(7, 9, 2),
207            InputEdge::new(9, 10, 2),
208            InputEdge::new(10, 8, 2),
209            InputEdge::new(8, 11, 2),
210            InputEdge::new(11, 6, 2),
211        ];
212        let graph = Graph::new(edges);
213
214        let mut tarjan = Tarjan::new();
215        assert_eq!(vec![6, 5, 3, 4, 5, 2, 1, 1, 1, 1, 1, 1], tarjan.run(&graph));
216    }
217}