competitive/graph/
sparse_graph.rs

1use super::{DirectedGraph, EdgeMap, Graph, MarkedScan, Neighbor, Scan, ScanSource, VertexMap};
2use std::{iter::Copied, marker::PhantomData, ops, slice};
3
4type Marker<T> = PhantomData<fn() -> T>;
5#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
6pub enum DirectedEdge {}
7#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
8pub enum UndirectedEdge {}
9#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
10pub enum BidirectionalEdge {}
11
12/// Static Sparse Graph represented as Compressed Sparse Row.
13#[derive(Debug, Clone)]
14pub struct SparseGraph<D> {
15    vsize: usize,
16    start: Vec<usize>,
17    neighbors: Vec<Neighbor<usize, usize>>,
18    pub edges: Vec<(usize, usize)>,
19    _marker: Marker<D>,
20}
21
22impl<D> SparseGraph<D> {
23    /// Return the number of vertices.
24    #[inline]
25    pub fn vertices_size(&self) -> usize {
26        self.vsize
27    }
28    /// Return the number of edges.
29    #[inline]
30    pub fn edges_size(&self) -> usize {
31        self.edges.len()
32    }
33    /// Return an iterator over graph vertices.
34    #[inline]
35    pub fn vertices(&self) -> ops::Range<usize> {
36        0..self.vertices_size()
37    }
38    pub fn builder<T>(vsize: usize) -> SparseGraphBuilder<T, D> {
39        SparseGraphBuilder::new(vsize)
40    }
41    pub fn builder_with_esize<T>(vsize: usize, esize: usize) -> SparseGraphBuilder<T, D> {
42        SparseGraphBuilder::new_with_esize(vsize, esize)
43    }
44}
45
46pub trait SparseGraphConstruction: Sized {
47    fn construct_graph(vsize: usize, edges: Vec<(usize, usize)>) -> SparseGraph<Self>;
48}
49
50impl<D> SparseGraph<D>
51where
52    D: SparseGraphConstruction,
53{
54    /// Construct graph from edges.
55    pub fn from_edges(vsize: usize, edges: Vec<(usize, usize)>) -> Self {
56        D::construct_graph(vsize, edges)
57    }
58    pub fn reverse_graph(&self) -> SparseGraph<D> {
59        let edges = self.edges.iter().map(|&(from, to)| (to, from)).collect();
60        D::construct_graph(self.vsize, edges)
61    }
62}
63
64impl SparseGraphConstruction for DirectedEdge {
65    fn construct_graph(vsize: usize, edges: Vec<(usize, usize)>) -> SparseGraph<Self> {
66        let mut start: Vec<_> = vec![0usize; vsize + 1];
67        for (from, _) in edges.iter().cloned() {
68            start[from] += 1;
69        }
70        for i in 1..=vsize {
71            start[i] += start[i - 1];
72        }
73        let mut neighbors = Vec::<Neighbor<usize, usize>>::with_capacity(edges.len());
74        let ptr = neighbors.as_mut_ptr();
75        for (id, (from, to)) in edges.iter().cloned().enumerate() {
76            start[from] -= 1;
77            unsafe { ptr.add(start[from]).write(Neighbor::new(to, id)) };
78        }
79        unsafe { neighbors.set_len(edges.len()) };
80        SparseGraph {
81            vsize,
82            start,
83            neighbors,
84            edges,
85            _marker: PhantomData,
86        }
87    }
88}
89
90impl SparseGraphConstruction for UndirectedEdge {
91    fn construct_graph(vsize: usize, edges: Vec<(usize, usize)>) -> SparseGraph<Self> {
92        let mut start: Vec<_> = vec![0usize; vsize + 1];
93        for (from, to) in edges.iter().cloned() {
94            start[to] += 1;
95            start[from] += 1;
96        }
97        for i in 1..=vsize {
98            start[i] += start[i - 1];
99        }
100        let mut neighbors = Vec::<Neighbor<usize, usize>>::with_capacity(edges.len() * 2);
101        let ptr = neighbors.as_mut_ptr();
102        for (id, (from, to)) in edges.iter().cloned().enumerate() {
103            start[from] -= 1;
104            unsafe { ptr.add(start[from]).write(Neighbor::new(to, id)) };
105            start[to] -= 1;
106            unsafe { ptr.add(start[to]).write(Neighbor::new(from, id)) };
107        }
108        unsafe { neighbors.set_len(edges.len() * 2) };
109        SparseGraph {
110            vsize,
111            start,
112            neighbors,
113            edges,
114            _marker: PhantomData,
115        }
116    }
117}
118
119impl SparseGraphConstruction for BidirectionalEdge {
120    fn construct_graph(vsize: usize, edges: Vec<(usize, usize)>) -> SparseGraph<Self> {
121        let mut start: Vec<_> = vec![0usize; vsize + 1];
122        for (from, to) in edges.iter().cloned() {
123            start[to] += 1;
124            start[from] += 1;
125        }
126        for i in 1..=vsize {
127            start[i] += start[i - 1];
128        }
129        let mut neighbors = Vec::<Neighbor<usize, usize>>::with_capacity(edges.len() * 2);
130        let ptr = neighbors.as_mut_ptr();
131        for (id, (from, to)) in edges.iter().cloned().enumerate() {
132            start[from] -= 1;
133            unsafe { ptr.add(start[from]).write(Neighbor::new(to, id * 2)) };
134            start[to] -= 1;
135            unsafe { ptr.add(start[to]).write(Neighbor::new(from, id * 2 + 1)) };
136        }
137        unsafe { neighbors.set_len(edges.len() * 2) };
138        SparseGraph {
139            vsize,
140            start,
141            neighbors,
142            edges,
143            _marker: PhantomData,
144        }
145    }
146}
147
148pub type DirectedSparseGraph = SparseGraph<DirectedEdge>;
149pub type UndirectedSparseGraph = SparseGraph<UndirectedEdge>;
150pub type BidirectionalSparseGraph = SparseGraph<BidirectionalEdge>;
151
152pub struct SparseGraphBuilder<T, D> {
153    vsize: usize,
154    edges: Vec<(usize, usize)>,
155    rest: Vec<T>,
156    _marker: Marker<D>,
157}
158impl<T, D> SparseGraphBuilder<T, D> {
159    pub fn new(vsize: usize) -> Self {
160        Self {
161            vsize,
162            edges: Default::default(),
163            rest: Default::default(),
164            _marker: PhantomData,
165        }
166    }
167    pub fn new_with_esize(vsize: usize, esize: usize) -> Self {
168        Self {
169            vsize,
170            edges: Vec::with_capacity(esize),
171            rest: Vec::with_capacity(esize),
172            _marker: PhantomData,
173        }
174    }
175    pub fn add_edge(&mut self, u: usize, v: usize, w: T) {
176        self.edges.push((u, v));
177        self.rest.push(w);
178    }
179}
180impl<T, D> SparseGraphBuilder<T, D>
181where
182    D: SparseGraphConstruction,
183{
184    pub fn build(self) -> (SparseGraph<D>, Vec<T>) {
185        let graph = SparseGraph::from_edges(self.vsize, self.edges);
186        (graph, self.rest)
187    }
188}
189
190pub struct SparseGraphScanner<U, T, D>
191where
192    U: Scan<Output = usize>,
193    T: Scan,
194{
195    vsize: usize,
196    esize: usize,
197    _marker: Marker<(U, T, D)>,
198}
199
200impl<U, T, D> SparseGraphScanner<U, T, D>
201where
202    U: Scan<Output = usize>,
203    T: Scan,
204{
205    pub fn new(vsize: usize, esize: usize) -> Self {
206        Self {
207            vsize,
208            esize,
209            _marker: PhantomData,
210        }
211    }
212}
213
214impl<U, T, D> MarkedScan for SparseGraphScanner<U, T, D>
215where
216    U: Scan<Output = usize>,
217    T: Scan,
218    D: SparseGraphConstruction,
219{
220    type Output = (SparseGraph<D>, Vec<<T as Scan>::Output>);
221    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
222        let mut builder = SparseGraphBuilder::new_with_esize(self.vsize, self.esize);
223        for _ in 0..self.esize {
224            builder.add_edge(U::scan(iter)?, U::scan(iter)?, T::scan(iter)?);
225        }
226        Some(builder.build())
227    }
228}
229
230pub type DirectedGraphScanner<U, T = ()> = SparseGraphScanner<U, T, DirectedEdge>;
231pub type UndirectedGraphScanner<U, T = ()> = SparseGraphScanner<U, T, UndirectedEdge>;
232pub type BidirectionalGraphScanner<U, T = ()> = SparseGraphScanner<U, T, BidirectionalEdge>;
233
234pub struct TreeGraphScanner<U, T = ()>
235where
236    U: Scan<Output = usize>,
237    T: Scan,
238{
239    vsize: usize,
240    _marker: Marker<(U, T)>,
241}
242impl<U, T> TreeGraphScanner<U, T>
243where
244    U: Scan<Output = usize>,
245    T: Scan,
246{
247    pub fn new(vsize: usize) -> Self {
248        Self {
249            vsize,
250            _marker: PhantomData,
251        }
252    }
253}
254impl<U, T> MarkedScan for TreeGraphScanner<U, T>
255where
256    U: Scan<Output = usize>,
257    T: Scan,
258{
259    type Output = (UndirectedSparseGraph, Vec<<T as Scan>::Output>);
260    fn mscan<I: ScanSource>(self, iter: &mut I) -> Option<Self::Output> {
261        UndirectedGraphScanner::<U, T>::new(self.vsize, self.vsize - 1).mscan(iter)
262    }
263}
264
265impl<D> Graph for SparseGraph<D>
266where
267    D: SparseGraphConstruction,
268{
269    type Vertex = usize;
270    type Label = usize;
271    type Vertices<'g>
272        = ops::Range<usize>
273    where
274        D: 'g;
275    type Neighbors<'g>
276        = Copied<slice::Iter<'g, Neighbor<usize, usize>>>
277    where
278        D: 'g;
279
280    #[inline]
281    fn vsize(&self) -> usize {
282        self.vsize
283    }
284
285    #[inline]
286    fn vertices(&self) -> Self::Vertices<'_> {
287        0..self.vsize
288    }
289
290    #[inline]
291    fn neighbors(&self, vertex: Self::Vertex) -> Self::Neighbors<'_> {
292        self.neighbors[self.start[vertex]..self.start[vertex + 1]]
293            .iter()
294            .copied()
295    }
296}
297
298impl DirectedGraph for SparseGraph<DirectedEdge> {}
299
300impl<T> EdgeMap<T> for SparseGraph<DirectedEdge> {
301    type Emap = Vec<T>;
302
303    #[inline]
304    fn construct_emap<F>(&self, f: F) -> Self::Emap
305    where
306        F: FnMut() -> T,
307    {
308        let mut map = Vec::with_capacity(self.edges.len());
309        map.resize_with(self.edges.len(), f);
310        map
311    }
312
313    #[inline]
314    fn emap_get<'a>(&self, map: &'a Self::Emap, eid: Self::Label) -> &'a T {
315        &map[eid]
316    }
317
318    #[inline]
319    fn emap_get_mut<'a>(&self, map: &'a mut Self::Emap, eid: Self::Label) -> &'a mut T {
320        &mut map[eid]
321    }
322}
323
324impl<T> EdgeMap<T> for SparseGraph<UndirectedEdge> {
325    type Emap = Vec<T>;
326
327    #[inline]
328    fn construct_emap<F>(&self, f: F) -> Self::Emap
329    where
330        F: FnMut() -> T,
331    {
332        let mut map = Vec::with_capacity(self.edges.len());
333        map.resize_with(self.edges.len(), f);
334        map
335    }
336
337    #[inline]
338    fn emap_get<'a>(&self, map: &'a Self::Emap, eid: Self::Label) -> &'a T {
339        &map[eid]
340    }
341
342    #[inline]
343    fn emap_get_mut<'a>(&self, map: &'a mut Self::Emap, eid: Self::Label) -> &'a mut T {
344        &mut map[eid]
345    }
346}
347
348impl<D, T> VertexMap<T> for SparseGraph<D>
349where
350    D: SparseGraphConstruction,
351{
352    type Vmap = Vec<T>;
353    #[inline]
354    fn construct_vmap<F>(&self, f: F) -> Self::Vmap
355    where
356        F: FnMut() -> T,
357    {
358        let mut v = Vec::with_capacity(self.vsize);
359        v.resize_with(self.vsize, f);
360        v
361    }
362    #[inline]
363    fn vmap_get<'a>(&self, map: &'a Self::Vmap, vid: Self::Vertex) -> &'a T {
364        &map[vid]
365    }
366    #[inline]
367    fn vmap_get_mut<'a>(&self, map: &'a mut Self::Vmap, vid: Self::Vertex) -> &'a mut T {
368        &mut map[vid]
369    }
370}
371
372#[cfg(test)]
373mod tests {
374    use super::*;
375    use crate::{rand, tools::Xorshift};
376
377    #[test]
378    fn test_sparse_graph() {
379        const Q: usize = 1_000;
380        const N: usize = 8;
381        const M: usize = 20;
382        let mut rng = Xorshift::default();
383        for _ in 0..Q {
384            rand!(rng, n: 1..=N, m: 0..=M, edges: [(0..n, 0..n); m]);
385
386            let directed = DirectedSparseGraph::from_edges(n, edges.clone());
387            let mut occurrences = vec![0; m];
388            for u in directed.vertices() {
389                for neighbor in directed.neighbors(u) {
390                    assert_eq!(edges[neighbor.label], (u, neighbor.to));
391                    occurrences[neighbor.label] += 1;
392                }
393            }
394            assert!(occurrences.into_iter().all(|count| count == 1));
395
396            let undirected = UndirectedSparseGraph::from_edges(n, edges.clone());
397            let mut occurrences = vec![0; m];
398            for u in undirected.vertices() {
399                for neighbor in undirected.neighbors(u) {
400                    let (from, to) = edges[neighbor.label];
401                    assert!((u == from && neighbor.to == to) || (u == to && neighbor.to == from));
402                    occurrences[neighbor.label] += 1;
403                }
404            }
405            assert!(occurrences.into_iter().all(|count| count == 2));
406
407            let bidirectional = BidirectionalSparseGraph::from_edges(n, edges.clone());
408            let mut occurrences = vec![0; m * 2];
409            for u in bidirectional.vertices() {
410                for neighbor in bidirectional.neighbors(u) {
411                    let (from, to) = edges[neighbor.label / 2];
412                    let expected = if neighbor.label % 2 == 0 {
413                        (from, to)
414                    } else {
415                        (to, from)
416                    };
417                    assert_eq!((u, neighbor.to), expected);
418                    occurrences[neighbor.label] += 1;
419                }
420            }
421            assert!(occurrences.into_iter().all(|count| count == 1));
422        }
423    }
424}
competitive/graph/sparse_graph.rs

competitive/graph/
sparse_graph.rs
