Skip to main content

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    use crate::graph::Graph;
99    use crate::tools::testutil::sample_usize;
100
101    fn is_connected(g: &UndirectedSparseGraph) -> bool {
102        let n = g.vertices_size();
103        if n == 0 {
104            return true;
105        }
106        let mut vis = vec![false; n];
107        let mut stack = vec![0];
108        let mut acc = 0usize;
109        vis[0] = true;
110        while let Some(u) = stack.pop() {
111            acc += 1;
112            for a in g.neighbors(u) {
113                if !vis[a.to] {
114                    vis[a.to] = true;
115                    stack.push(a.to);
116                }
117            }
118        }
119        acc == n
120    }
121
122    fn is_tree(g: &UndirectedSparseGraph) -> bool {
123        let n = g.vertices_size();
124        let m = g.edges_size();
125        (n == 0 || m + 1 == n) && is_connected(g)
126    }
127
128    #[test]
129    fn test_tree_generators() {
130        let mut rng = Xorshift::default();
131        for n in (0..=20).chain(sample_usize(&mut rng, 16, 0..=10_000, 100)) {
132            for graph in [
133                rng.random(PruferSequence(n)),
134                rng.random(MixedTree(n)),
135                rng.random(PathTree(n)),
136                rng.random(StarTree(n)),
137            ] {
138                assert_eq!(graph.vertices_size(), n);
139                assert!(is_tree(&graph));
140            }
141            let path = rng.random(PathTree(n));
142            assert!(path.vertices().all(|v| path.neighbors(v).len() <= 2));
143            let star = rng.random(StarTree(n));
144            if n > 0 {
145                assert_eq!(star.neighbors(0).len(), n - 1);
146            }
147        }
148    }
149}