competitive_programming_rs/data_structure/
persistent_array.rs

1pub mod persistent_array {
2    use std::rc::Rc;
3
4    const N: usize = 20;
5    #[derive(Clone)]
6    pub struct Node<T: Clone> {
7        data: Option<T>,
8        children: [Option<Rc<Node<T>>>; N],
9    }
10
11    impl<T: Clone> Node<T> {
12        pub fn new(data: Option<T>, children: [Option<Rc<Node<T>>>; N]) -> Self {
13            Self { data, children }
14        }
15    }
16
17    impl<T: Clone> Default for Node<T> {
18        fn default() -> Self {
19            Self {
20                data: None,
21                children: Default::default(),
22            }
23        }
24    }
25
26    pub fn set<T: Clone>(index: usize, data: T, node: Option<Rc<Node<T>>>) -> Rc<Node<T>> {
27        if index == 0 {
28            match node {
29                Some(node) => {
30                    let new_node = Node::new(Some(data), node.children.clone());
31                    Rc::new(new_node)
32                }
33                None => Rc::new(Node::new(Some(data), Default::default())),
34            }
35        } else {
36            let child = match node
37                .as_ref()
38                .and_then::<&Rc<Node<T>>, _>(|node| node.children[index % N].as_ref())
39            {
40                Some(next_node) => set(index / N, data, Some(next_node.clone())),
41                None => set(index / N, data, None),
42            };
43            let mut new_node = match node {
44                Some(node) => node.as_ref().clone(),
45                None => Node::default(),
46            };
47            new_node.children[index % N] = Some(child);
48            Rc::new(new_node)
49        }
50    }
51
52    pub fn get<T: Clone>(index: usize, node: &Rc<Node<T>>) -> Option<T> {
53        if index == 0 {
54            node.data.clone()
55        } else {
56            match node.children[index % N].as_ref() {
57                Some(next_node) => get(index / N, next_node),
58                None => None,
59            }
60        }
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::persistent_array::{get, set, Node};
67    use rand::distributions::Uniform;
68    use rand::{thread_rng, Rng};
69    use std::rc::Rc;
70
71    #[test]
72    fn test_persistent_array() {
73        let mut rng = thread_rng();
74        let n: usize = 20;
75        let mut vs = vec![vec![None; n]];
76        let mut nodes = vec![Rc::new(Node::default())];
77
78        for _ in 0..10000 {
79            let pick = rng.sample(Uniform::from(0..vs.len()));
80            let mut new_vec = vs[pick].clone();
81            let node = nodes[pick].clone();
82            let pos = rng.sample(Uniform::from(0..n));
83            let value: i64 = rng.gen();
84            new_vec[pos] = Some(value);
85            let new_node = set(pos, value, Some(node));
86
87            for i in 0..n {
88                let expected = new_vec[i];
89                let actual = get(i, &new_node);
90                assert_eq!(expected, actual);
91            }
92            vs.push(new_vec.clone());
93            nodes.push(new_node.clone());
94
95            let value: i64 = rng.gen();
96            new_vec[pos] = Some(value);
97            let new_node = set(pos, value, Some(new_node));
98            for i in 0..n {
99                let expected = new_vec[i];
100                let actual = get(i, &new_node);
101                assert_eq!(expected, actual);
102            }
103            vs.push(new_vec.clone());
104            nodes.push(new_node);
105        }
106    }
107}
competitive_programming_rs/data_structure/persistent_array.rs

competitive_programming_rs/data_structure/
persistent_array.rs
