rustgym_util/
graph.rs

1use std::collections::HashSet;
2
3#[derive(Debug, Eq, PartialEq)]
4struct Node {
5    id: usize,
6    discovery: usize,
7    lowest: usize,
8    on_stack: bool,
9}
10
11impl Node {
12    pub fn new(id: usize) -> Self {
13        let discovery = std::usize::MAX;
14        let lowest = std::usize::MAX;
15        let on_stack = false;
16        Node {
17            id,
18            discovery,
19            lowest,
20            on_stack,
21        }
22    }
23
24    fn visit(&mut self, time: usize) {
25        self.discovery = time;
26    }
27
28    fn is_visited(&self) -> bool {
29        self.discovery != std::usize::MAX
30    }
31}
32
33#[derive(Debug, Eq, PartialEq)]
34pub enum EdgeKind {
35    TreeEdge,
36    BackEdge,
37    CrossEdge,
38    ForwardEdge,
39}
40
41#[derive(Debug, Eq, PartialEq)]
42pub struct Edge {
43    pub tail: usize,
44    pub head: usize,
45    pub kind: EdgeKind,
46}
47
48impl Edge {
49    fn new(tail: usize, head: usize, kind: EdgeKind) -> Self {
50        Edge { tail, head, kind }
51    }
52}
53
54#[derive(Debug, Eq, PartialEq)]
55pub struct Graph {
56    pub n: usize,
57    pub edges: Vec<Edge>,
58    adj: Vec<HashSet<usize>>,
59    nodes: Vec<Node>,
60    stack: Vec<usize>,
61    time: usize,
62}
63
64impl Graph {
65    pub fn new(n: usize) -> Self {
66        let time = 0;
67        let adj = vec![HashSet::new(); n];
68        let nodes = (0..n).map(Node::new).collect();
69        let edges = vec![];
70        let stack = vec![];
71        Graph {
72            n,
73            time,
74            adj,
75            nodes,
76            edges,
77            stack,
78        }
79    }
80
81    pub fn init_with_edges(&mut self, edges: Vec<Vec<i32>>, directed: bool) {
82        for edge in edges {
83            let u = edge[0] as usize;
84            let v = edge[1] as usize;
85            if !directed {
86                self.adj[u].insert(v);
87                self.adj[v].insert(u);
88            } else {
89                self.adj[u].insert(v);
90            }
91        }
92    }
93
94    pub fn dfs(&mut self, u: usize, parent: usize) {
95        self.nodes[u].visit(self.time);
96        self.time += 1;
97        let neighbours: Vec<usize> = self.adj[u].iter().copied().collect();
98        for v in neighbours {
99            if v == parent {
100                continue;
101            }
102            if !self.nodes[v].is_visited() {
103                self.edges.push(Edge::new(u, v, EdgeKind::TreeEdge));
104                self.dfs(v, u);
105            } else {
106                self.edges.push(Edge::new(u, v, EdgeKind::BackEdge));
107            }
108        }
109    }
110
111    pub fn travase(&mut self) {
112        for u in 0..self.n {
113            if !self.nodes[u].is_visited() {
114                self.dfs(u, u);
115            }
116        }
117    }
118}