okvs/
graph.rs

1use std::collections::{hash_map::Entry, HashMap};
2
3/// An edge with a value.
4#[derive(Debug, PartialEq, Clone, Copy)]
5pub struct Edge<V> {
6    pub to_vertex: usize,
7    pub value: V,
8}
9
10/// An undirected graph with a value on each edge.
11#[derive(Debug, PartialEq)]
12pub struct UndirectedGraph<V> {
13    pub(crate) adjacency_list: Vec<Vec<Edge<V>>>,
14    pub(crate) edge_count: usize,
15    self_references: HashMap<usize, Edge<V>>,
16}
17
18impl<V: Copy> UndirectedGraph<V> {
19    pub fn new(vertex_count: usize) -> Self {
20        Self {
21            adjacency_list: (0..vertex_count).map(|_| vec![]).collect(),
22            edge_count: 0,
23            self_references: HashMap::new(),
24        }
25    }
26
27    pub fn has_edges(&self, vertex: usize) -> bool {
28        if self.self_references.contains_key(&vertex) {
29            return true;
30        }
31
32        !self.adjacency_list[vertex].is_empty()
33    }
34
35    /// Adds an edge between `vertex_a` and `vertex_b` with the specified `value`.
36    /// Returns `true` if successful, and `false` if an edge already existed with a different value.
37    pub fn add_edge(&mut self, vertex_a: usize, vertex_b: usize, value: V) -> bool {
38        for edge in &self.adjacency_list[vertex_a] {
39            if edge.to_vertex == vertex_b {
40                return false;
41            }
42        }
43
44        if vertex_a == vertex_b {
45            if let Entry::Vacant(e) = self.self_references.entry(vertex_a) {
46                e.insert(Edge {
47                    to_vertex: vertex_b,
48                    value,
49                });
50                self.edge_count += 1;
51                return true;
52            } else {
53                return false;
54            }
55        }
56
57        self.adjacency_list[vertex_a].push(Edge {
58            to_vertex: vertex_b,
59            value,
60        });
61        self.adjacency_list[vertex_b].push(Edge {
62            to_vertex: vertex_a,
63            value,
64        });
65        self.edge_count += 1;
66        true
67    }
68
69    pub fn pop_edge(&mut self, vertex: usize) -> Option<Edge<V>> {
70        if self.self_references.contains_key(&vertex) {
71            self.edge_count -= 1;
72            return self.self_references.remove(&vertex);
73        }
74
75        let edge = self.adjacency_list[vertex].pop()?;
76        let index = self.adjacency_list[edge.to_vertex]
77            .iter()
78            .position(|edge| edge.to_vertex == vertex)
79            .unwrap();
80        self.adjacency_list[edge.to_vertex].swap_remove(index);
81        self.edge_count -= 1;
82
83        Some(edge)
84    }
85
86    /// Returns the edge if it is the only one remaining. A self-reference does not count, so the function returns `None`.
87    pub fn pop_only_edge(&mut self, vertex: usize) -> Option<Edge<V>> {
88        if self.adjacency_list[vertex].len() != 1 {
89            return None;
90        }
91
92        if self.self_references.contains_key(&vertex) {
93            return None;
94        }
95
96        let edge = self.adjacency_list[vertex].pop()?;
97        let index = self.adjacency_list[edge.to_vertex]
98            .iter()
99            .position(|edge| edge.to_vertex == vertex)
100            .unwrap();
101        self.adjacency_list[edge.to_vertex].swap_remove(index);
102        self.edge_count -= 1;
103
104        Some(edge)
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::UndirectedGraph;
111
112    // TODO: Add testcases for self-references
113
114    #[test]
115    fn empty_graph() {
116        let mut graph: UndirectedGraph<usize> = UndirectedGraph::new(10);
117        assert_eq!(graph.edge_count, 0);
118        for i in 0..10 {
119            assert!(graph.pop_only_edge(i).is_none());
120        }
121    }
122
123    #[test]
124    fn add_edge() {
125        let mut graph: UndirectedGraph<usize> = UndirectedGraph::new(10);
126        graph.add_edge(5, 1, 2500);
127        assert_eq!(graph.edge_count, 1);
128    }
129
130    #[test]
131    fn add_edge_invariant() {
132        let mut graph_a: UndirectedGraph<usize> = UndirectedGraph::new(10);
133        graph_a.add_edge(5, 1, 3000);
134
135        let mut graph_b: UndirectedGraph<usize> = UndirectedGraph::new(10);
136        graph_b.add_edge(1, 5, 3000);
137
138        assert_eq!(graph_a, graph_b);
139    }
140
141    #[test]
142    fn pop_edge() {
143        let mut graph: UndirectedGraph<usize> = UndirectedGraph::new(10);
144        graph.add_edge(5, 1, 2500);
145
146        let edge_a = graph.pop_edge(1).unwrap();
147
148        assert_eq!(graph.edge_count, 0);
149        assert_eq!(edge_a.value, 2500);
150        assert_eq!(edge_a.to_vertex, 5);
151        assert!(graph.adjacency_list[1].is_empty());
152        assert!(graph.adjacency_list[5].is_empty());
153
154        graph.add_edge(5, 1, 2500);
155        assert_eq!(graph.edge_count, 1);
156
157        let edge_b = graph.pop_edge(5).unwrap();
158
159        assert_eq!(graph.edge_count, 0);
160        assert_eq!(edge_b.value, 2500);
161        assert_eq!(edge_b.to_vertex, 1);
162        assert!(graph.adjacency_list[1].is_empty());
163        assert!(graph.adjacency_list[5].is_empty());
164    }
165
166    #[test]
167    fn pop_only_edge() {
168        let mut graph: UndirectedGraph<usize> = UndirectedGraph::new(10);
169        graph.add_edge(5, 1, 2500);
170
171        let edge_a = graph.pop_only_edge(1).unwrap();
172
173        assert_eq!(graph.edge_count, 0);
174        assert_eq!(edge_a.value, 2500);
175        assert_eq!(edge_a.to_vertex, 5);
176        assert!(graph.adjacency_list[1].is_empty());
177        assert!(graph.adjacency_list[5].is_empty());
178
179        graph.add_edge(5, 1, 2500);
180        assert_eq!(graph.edge_count, 1);
181
182        let edge_b = graph.pop_only_edge(5).unwrap();
183
184        assert_eq!(graph.edge_count, 0);
185        assert_eq!(edge_b.value, 2500);
186        assert_eq!(edge_b.to_vertex, 1);
187        assert!(graph.adjacency_list[1].is_empty());
188        assert!(graph.adjacency_list[5].is_empty());
189    }
190}
okvs/graph.rs

okvs/
graph.rs
