competitive/tree/
generator.rs

1use super::{RandomSpec, UndirectedSparseGraph, Xorshift};
2
3/// Generate Tree with Prüfer sequence
4pub struct PruferSequence<T>(pub T);
5
6impl<T: RandomSpec<usize>> RandomSpec<UndirectedSparseGraph> for PruferSequence<T> {
7    fn rand(&self, rng: &mut Xorshift) -> UndirectedSparseGraph {
8        let n = rng.random(&self.0);
9        let edges = from_prufer_sequence(
10            n,
11            &rng.random_iter(0..n)
12                .take(n.saturating_sub(2))
13                .collect::<Vec<usize>>(),
14        );
15        UndirectedSparseGraph::from_edges(n, edges)
16    }
17}
18
19pub struct PathTree<T>(pub T);
20
21impl<T: RandomSpec<usize>> RandomSpec<UndirectedSparseGraph> for PathTree<T> {
22    fn rand(&self, rng: &mut Xorshift) -> UndirectedSparseGraph {
23        let n = rng.random(&self.0);
24        let edges = (1..n).map(|u| (u - 1, u)).collect();
25        UndirectedSparseGraph::from_edges(n, edges)
26    }
27}
28
29pub struct StarTree<T>(pub T);
30
31impl<T: RandomSpec<usize>> RandomSpec<UndirectedSparseGraph> for StarTree<T> {
32    fn rand(&self, rng: &mut Xorshift) -> UndirectedSparseGraph {
33        let n = rng.random(&self.0);
34        let edges = (1..n).map(|u| (0, u)).collect();
35        UndirectedSparseGraph::from_edges(n, edges)
36    }
37}
38
39pub struct MixedTree<T>(pub T);
40
41impl<T: RandomSpec<usize>> RandomSpec<UndirectedSparseGraph> for MixedTree<T> {
42    fn rand(&self, rng: &mut Xorshift) -> UndirectedSparseGraph {
43        fn rand_inner(n: usize, rng: &mut Xorshift) -> Vec<(usize, usize)> {
44            let mut edges = Vec::with_capacity(n.saturating_sub(1));
45            if n >= 2 {
46                let k = rng.random(1..n);
47                for n in [k, n - k].iter().cloned() {
48                    let ty = rng.rand(6);
49                    edges.extend(match ty {
50                        0 => from_prufer_sequence(
51                            n,
52                            &rng.random_iter(0..n)
53                                .take(n.saturating_sub(2))
54                                .collect::<Vec<usize>>(),
55                        ),
56                        1 => (1..n).map(|u| (u - 1, u)).collect(),
57                        2 => (1..n).map(|u| (0, u)).collect(),
58                        _ => rand_inner(n, rng),
59                    });
60                }
61                for (u, v) in edges[k - 1..].iter_mut() {
62                    *u += k;
63                    *v += k;
64                }
65                edges.push((rng.random(0..k), rng.random(k..n)));
66            }
67            edges
68        }
69        let n = rng.random(&self.0);
70        let edges = rand_inner(n, rng);
71        UndirectedSparseGraph::from_edges(n, edges)
72    }
73}
74
75fn from_prufer_sequence(n: usize, prufer: &[usize]) -> Vec<(usize, usize)> {
76    use std::collections::BinaryHeap;
77    let mut edges = Vec::with_capacity(n.saturating_sub(1));
78    if n >= 2 {
79        let mut deg = vec![0usize; n];
80        prufer.iter().for_each(|&a| deg[a] += 1);
81        let mut heap: BinaryHeap<usize> = (0..n).filter(|&u| deg[u] == 0).collect();
82        for &a in prufer {
83            deg[a] -= 1;
84            let b = heap.pop().unwrap();
85            edges.push((a, b));
86            if deg[a] == 0 {
87                heap.push(a);
88            }
89        }
90        edges.push((heap.pop().unwrap(), heap.pop().unwrap()));
91    }
92    edges
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    fn is_connected(g: &UndirectedSparseGraph) -> bool {
100        let n = g.vertices_size();
101        if n == 0 {
102            return true;
103        }
104        let mut vis = vec![false; n];
105        let mut stack = vec![0];
106        let mut acc = 0usize;
107        vis[0] = true;
108        while let Some(u) = stack.pop() {
109            acc += 1;
110            for a in g.adjacencies(u) {
111                if !vis[a.to] {
112                    vis[a.to] = true;
113                    stack.push(a.to);
114                }
115            }
116        }
117        acc == n
118    }
119
120    fn is_tree(g: &UndirectedSparseGraph) -> bool {
121        let n = g.vertices_size();
122        let m = g.edges_size();
123        (n == 0 || m + 1 == n) && is_connected(g)
124    }
125
126    #[test]
127    fn prufer_sequence_small() {
128        const Q: usize = 10_000;
129        const N: usize = 20;
130        let mut rng = Xorshift::default();
131        for _ in 0..Q {
132            let g = rng.random(PruferSequence(1..=N));
133            assert!(is_tree(&g));
134        }
135    }
136
137    #[test]
138    fn prufer_sequence_big() {
139        const Q: usize = 20;
140        const N: usize = 10_000;
141        let mut rng = Xorshift::default();
142        for _ in 0..Q {
143            let g = rng.random(PruferSequence(N - Q..=N));
144            assert!(is_tree(&g));
145        }
146    }
147
148    #[test]
149    fn path_small() {
150        const N: usize = 20;
151        let mut rng = Xorshift::default();
152        for n in 0..=N {
153            let g = rng.random(PathTree(n));
154            assert!(is_tree(&g));
155        }
156    }
157
158    #[test]
159    fn path_big() {
160        const Q: usize = 20;
161        const N: usize = 10_000;
162        let mut rng = Xorshift::default();
163        for n in N - Q..=N {
164            let g = rng.random(PathTree(n));
165            assert!(is_tree(&g));
166        }
167    }
168
169    #[test]
170    fn star_small() {
171        const N: usize = 20;
172        let mut rng = Xorshift::default();
173        for n in 0..=N {
174            let g = rng.random(StarTree(n));
175            assert!(is_tree(&g));
176        }
177    }
178
179    #[test]
180    fn star_big() {
181        const Q: usize = 20;
182        const N: usize = 10_000;
183        let mut rng = Xorshift::default();
184        for n in N - Q..=N {
185            let g = rng.random(StarTree(n));
186            assert!(is_tree(&g));
187        }
188    }
189
190    #[test]
191    fn mixed_small() {
192        const Q: usize = 10_000;
193        const N: usize = 20;
194        let mut rng = Xorshift::default();
195        for _ in 0..Q {
196            let g = rng.random(MixedTree(1..=N));
197            assert!(is_tree(&g));
198        }
199    }
200
201    #[test]
202    fn mixed_big() {
203        const Q: usize = 20;
204        const N: usize = 10_000;
205        let mut rng = Xorshift::default();
206        for _ in 0..Q {
207            let g = rng.random(MixedTree(N - Q..=N));
208            assert!(is_tree(&g));
209        }
210    }
211}
competitive/tree/generator.rs

competitive/tree/
generator.rs
