placements_tree/
lib.rs

1mod apply;
2mod fill;
3mod max;
4mod node;
5mod recalc;
6
7pub use crate::apply::Apply;
8use crate::fill::Fill;
9pub use crate::max::Max;
10use crate::node::Node;
11pub use crate::recalc::Recalc;
12use std::collections::LinkedList;
13use std::ptr::NonNull;
14
15pub struct PlacementsTree<V, E, D> {
16    _root: Box<Node<D>>,
17    vertices: Vec<V>,
18    vertices_idx: Vec<LinkedList<NonNull<Node<D>>>>,
19    edges: Vec<Vec<E>>,
20    edges_idx: Vec<Vec<LinkedList<NonNull<Node<D>>>>>,
21    n: usize,
22}
23
24impl<V, E, D> PlacementsTree<V, E, D> {
25    pub fn new(n: usize, k: usize, key: usize, val: D) -> Self
26    where
27        V: Default + Clone,
28        E: Default + Clone,
29        D: Max,
30    {
31        assert!(key <= n);
32        let k = k.min(n);
33        let root = Node::root(n, k, key, val);
34        let vertices = vec![V::default(); n + 1];
35        let mut vertices_idx = vec![LinkedList::new(); n + 1];
36        root.fill(&mut vertices_idx);
37        let edges = vec![vec![E::default(); n + 1]; n + 1];
38        let mut edges_idx = vec![vec![LinkedList::new(); n + 1]; n + 1];
39        root.fill(&mut edges_idx);
40        Self {
41            _root: root,
42            vertices,
43            vertices_idx,
44            edges,
45            edges_idx,
46            n,
47        }
48    }
49
50    pub fn update_vertex<Diff>(&mut self, v: usize, diff: Diff) -> Option<&D>
51    where
52        V: Apply<Diff>,
53        D: Recalc<V, E> + PartialOrd,
54    {
55        assert!(v <= self.n);
56        self.vertices[v].apply(diff);
57        if self.vertices_idx[v].is_empty() {
58            None
59        } else {
60            unsafe {
61                let mut vertices = self.vertices_idx[v].iter_mut();
62                let mut shortest = vertices
63                    .next()
64                    .map(|vertex| vertex.as_mut().recalc_children(&self.vertices, &self.edges))
65                    .unwrap();
66                for vertex in vertices {
67                    let recalced = vertex.as_mut().recalc_children(&self.vertices, &self.edges);
68                    if recalced < shortest {
69                        shortest = recalced;
70                    }
71                }
72                Some(shortest)
73            }
74        }
75    }
76
77    pub fn update_edge<Diff>(&mut self, v: usize, u: usize, diff: Diff) -> Option<&D>
78    where
79        E: Apply<Diff>,
80        D: Recalc<V, E> + PartialOrd,
81    {
82        assert!(v <= self.n);
83        assert!(u <= self.n);
84        assert!(v != u);
85        self.edges[v][u].apply(diff);
86        if self.edges_idx[v][u].is_empty() {
87            None
88        } else {
89            unsafe {
90                let mut edges = self.edges_idx[v][u].iter_mut();
91                let mut shortest = edges
92                    .next()
93                    .map(|edge| edge.as_mut().recalc(&self.vertices, &self.edges))
94                    .unwrap();
95                for edge in edges {
96                    let recalced = edge.as_mut().recalc(&self.vertices, &self.edges);
97                    if recalced < shortest {
98                        shortest = recalced;
99                    }
100                }
101                Some(shortest)
102            }
103        }
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[derive(PartialEq, Eq, PartialOrd, Debug)]
112    struct Dist(i64);
113
114    impl Max for Dist {
115        fn max() -> Self {
116            Dist(i64::MAX)
117        }
118    }
119
120    impl Recalc<i64, i64> for Dist {
121        fn recalc(&self, vertex: &i64, edge: &i64) -> Self {
122            if *self == Self::max() {
123                Self::max()
124            } else {
125                Self(self.0 + vertex + edge)
126            }
127        }
128    }
129
130    #[test]
131    #[should_panic(expected = "assertion failed: key <= n")]
132    fn new_panicked_test() {
133        PlacementsTree::<i64, i64, Dist>::new(2, 2, 3, Dist(0));
134    }
135
136    #[test]
137    fn update_test() {
138        let mut ptree: PlacementsTree<i64, i64, Dist> = PlacementsTree::new(2, 2, 0, Dist(0));
139        assert_eq!(*ptree.update_vertex(1, 1).unwrap(), Dist::max());
140        assert_eq!(*ptree.update_vertex(2, 1).unwrap(), Dist::max());
141        assert_eq!(*ptree.update_edge(0, 1, 1).unwrap(), Dist(3));
142        assert_eq!(*ptree.update_edge(0, 2, 2).unwrap(), Dist(4));
143        assert_eq!(*ptree.update_edge(1, 0, 3).unwrap(), Dist(7));
144        assert_eq!(*ptree.update_edge(1, 2, 4).unwrap(), Dist(7));
145        assert_eq!(*ptree.update_edge(2, 0, 5).unwrap(), Dist(12));
146        assert_eq!(*ptree.update_edge(2, 1, 6).unwrap(), Dist(13));
147        assert_eq!(*ptree.update_vertex(0, 1).unwrap(), Dist(13));
148    }
149
150    #[test]
151    #[should_panic(expected = "assertion failed: v <= self.n")]
152    fn update_vertex_panicked_test() {
153        let mut ptree: PlacementsTree<i64, i64, Dist> = PlacementsTree::new(2, 2, 0, Dist(0));
154        ptree.update_vertex(3, 0);
155    }
156
157    #[test]
158    #[should_panic(expected = "assertion failed: v <= self.n")]
159    fn update_edge_panicked_1_test() {
160        let mut ptree: PlacementsTree<i64, i64, Dist> = PlacementsTree::new(2, 2, 0, Dist(0));
161        ptree.update_edge(3, 0, 0);
162    }
163
164    #[test]
165    #[should_panic(expected = "assertion failed: u <= self.n")]
166    fn update_edge_panicked_2_test() {
167        let mut ptree: PlacementsTree<i64, i64, Dist> = PlacementsTree::new(2, 2, 0, Dist(0));
168        ptree.update_edge(0, 3, 0);
169    }
170
171    #[test]
172    #[should_panic(expected = "assertion failed: v != u")]
173    fn update_edge_panicked_3_test() {
174        let mut ptree: PlacementsTree<i64, i64, Dist> = PlacementsTree::new(2, 2, 0, Dist(0));
175        ptree.update_edge(0, 0, 0);
176    }
177}
placements_tree/lib.rs

placements_tree/
lib.rs
