graph_tools/
csr.rs

1pub struct CSR{
2    nodes: Vec<usize>,
3    edges: Vec<usize>
4}
5
6impl CSR{
7
8    pub fn from_sorted_edges(edges: &[(usize, usize)], n_nodes: usize) -> CSR{
9        let n_edges = edges.len();
10
11        
12        let mut csr = CSR{
13            edges: vec![0; n_edges],
14            nodes: vec![0; n_nodes + 1]
15        };
16        
17        for (i, (u, v)) in edges.iter().enumerate() {
18            csr.edges[i] = *v;
19            csr.nodes[*u + 1] += 1;
20        }
21        
22        for i in 1..n_nodes {
23            csr.nodes[i+1] += csr.nodes[i];
24        }
25
26        return csr;
27    }
28
29    pub fn from_edges(edges: &[(usize, usize)], n_nodes: usize) -> CSR{
30        let n_edges = edges.len();
31        let mut csr = CSR{
32            edges: vec![0; n_edges*2],
33            nodes: vec![0; n_nodes + 1],
34        };
35
36        for (u, v) in edges {
37            csr.nodes[*u+1] += 1;
38            csr.nodes[*v+1] += 1;
39        }
40
41        for i in 0..n_nodes {
42            csr.nodes[i+1] += csr.nodes[i];
43        }
44        
45        for (u, v) in edges {
46            csr.edges[csr.nodes[*u]] = *v;
47            csr.nodes[*u] += 1;
48            csr.edges[csr.nodes[*v]] = *u;
49            csr.nodes[*v] += 1;
50        }
51        
52        for i in (0..n_nodes).rev() {
53            csr.nodes[i+1] = csr.nodes[i];
54        }
55        csr.nodes[0] = 0;
56
57        return csr;
58    }
59
60    pub fn degree(&self, u: usize) -> usize {
61        return self.nodes[u+1] - self.nodes[u];
62    }
63
64    pub fn neighbors(&self, u: usize) -> &[usize] {
65        return &self.edges[self.nodes[u]..self.nodes[u+1]];
66    }
67
68    pub fn iter_edges(&self) -> CSREdgeIterator{
69        return CSREdgeIterator{
70            csr: self,
71            n_cur: 0,
72            e_cur: 0
73        }
74    }
75
76    pub fn n_nodes(&self) -> usize{
77        return self.nodes.len() - 1;
78    }
79}
80
81pub struct CSREdgeIterator<'a>{
82    csr: &'a CSR,
83    n_cur: usize,
84    e_cur: usize
85}
86
87impl<'a> Iterator for CSREdgeIterator<'a> {
88    type Item = (usize, usize);
89
90    fn next(&mut self) -> Option<Self::Item> {
91        
92        if self.e_cur < self.csr.edges.len() {
93            
94            while self.csr.nodes[self.n_cur+1] <= self.e_cur {
95                self.n_cur += 1;
96            }
97
98            self.e_cur += 1;
99
100            return Some((self.n_cur, self.csr.edges[self.e_cur-1]));
101        }
102
103        return None;
104    }
105}