Preliminary implementation of a generic ArqTree

Author: EbTech

src/arq_tree.rs

@@ -7,61 +7,62 @@
 ///
 /// - modify(l, r, f) replaces a_i (l <= i <= r) by f(a_i) for a homomorphism f
 /// - query(l, r) returns the aggregate a_l + a_{l+1} + ... + a_r
-///
-/// To customize, simply change the commented lines.
-/// In this example, we chose to support range sum queries and range constant
-/// assignments. Since constant assignment f_c(a) = c is not a homomorphism over
-/// integers, we have to augment the monoid type, using the 2D vector (a_i, 1)
-/// instead of a_i. You may check that f_c((a, s)) = (c*s, s) is a homomorphism.
-pub struct ArqTree {
-    d: Vec<Option<i64>>,
-    t: Vec<i64>,
-    s: Vec<i64>,
+pub struct ArqTree<T: ArqSpec> {
+    d: Vec<Option<T::H>>,
+    t: Vec<T::M>,
 }
 
-impl ArqTree {
-    /// Initializes a sequence of identity elements.
-    pub fn new(size: usize) -> Self {
-        let mut s = vec![1; 2 * size];
+impl<T: ArqSpec> ArqTree<T>
+where
+    T::H: Clone,
+{
+    /// Initializes a static balanced tree on top of the given sequence.
+    pub fn new(mut init: Vec<T::M>) -> Self {
+        let size = init.len();
+        let mut t = (0..size).map(|_| T::identity()).collect::<Vec<_>>();
+        t.append(&mut init);
         for i in (0..size).rev() {
-            s[i] = s[i << 1] + s[i << 1 | 1];
+            t[i] = T::op(&t[i << 1], &t[i << 1 | 1]);
         }
         Self {
-            d: vec![None; size],
-            t: vec![0; 2 * size], // monoid identity
-            s: s,
+            d: (0..size).map(|_| None).collect(),
+            t: t,
         }
     }
 
-    fn apply(&mut self, p: usize, f: i64) {
-        self.t[p] = f * self.s[p]; // hom application
+    fn apply(&mut self, p: usize, f: &T::H) {
+        self.t[p] = T::apply(f, &self.t[p]); // hom application
         if p < self.d.len() {
-            self.d[p] = Some(f); // hom composition
+            let h = if let Some(ref g) = self.d[p] {
+                T::compose(f, g)
+            } else {
+                f.clone()
+            };
+            self.d[p] = Some(h);
         }
     }
 
     fn push(&mut self, p: usize) {
         for s in (1..32).rev() {
             let i = p >> s;
-            if let Some(f) = self.d[i] {
+            if let Some(ref f) = self.d[i].take() {
                 self.apply(i << 1, f);
                 self.apply(i << 1 | 1, f);
-                self.d[i] = None;
             }
         }
     }
 
     fn pull(&mut self, mut p: usize) {
         while p > 1 {
             p >>= 1;
-            if self.d[p] == None {
-                self.t[p] = self.t[p << 1] + self.t[p << 1 | 1]; // monoid op
+            if self.d[p].is_none() {
+                self.t[p] = T::op(&self.t[p << 1], &self.t[p << 1 | 1]); // monoid op
             }
         }
     }
 
     /// Performs the homomorphism f on all entries from l to r, inclusive.
-    pub fn modify(&mut self, mut l: usize, mut r: usize, f: i64) {
+    pub fn modify(&mut self, mut l: usize, mut r: usize, f: &T::H) {
         l += self.d.len();
         r += self.d.len();
         let (l0, r0) = (l, r);
@@ -84,19 +85,19 @@ impl ArqTree {
     }
 
     /// Returns the aggregate range query on all entries from l to r, inclusive.
-    pub fn query(&mut self, mut l: usize, mut r: usize) -> i64 {
+    pub fn query(&mut self, mut l: usize, mut r: usize) -> T::M {
         l += self.d.len();
         r += self.d.len();
         self.push(l);
         self.push(r);
-        let mut res = 0; // monoid identity
+        let mut res = T::identity(); // monoid identity
         while l <= r {
             if l & 1 == 1 {
-                res = res + self.t[l]; // monoid op
+                res = T::op(&res, &self.t[l]); // monoid op
                 l += 1;
             }
             if r & 1 == 0 {
-                res = self.t[r] + res; // monoid op
+                res = T::op(&self.t[r], &res); // monoid op
                 r -= 1;
             }
             l >>= 1;
@@ -106,17 +107,52 @@ impl ArqTree {
     }
 }
 
+pub trait ArqSpec {
+    type H;
+    type M;
+    fn compose(f: &Self::H, g: &Self::H) -> Self::H;
+    fn apply(f: &Self::H, a: &Self::M) -> Self::M;
+    fn op(a: &Self::M, b: &Self::M) -> Self::M;
+    fn identity() -> Self::M;
+}
+
+/// In this example, we chose to support range sum queries and range constant
+/// assignments. Since constant assignment f_c(a) = c is not a homomorphism over
+/// integers, we have to augment the monoid type, using the 2D vector (a_i, 1)
+/// instead of a_i. You may check that f_c((a, s)) = (c*s, s) is a homomorphism.
+pub struct AssignAdd;
+
+impl ArqSpec for AssignAdd {
+    type H = i64;
+    type M = (i64, i64);
+    fn compose(f: &Self::H, _: &Self::H) -> Self::H {
+        *f
+    }
+    fn apply(f: &Self::H, a: &Self::M) -> Self::M {
+        (f * a.1, a.1)
+    }
+    fn op(a: &Self::M, b: &Self::M) -> Self::M {
+        (a.0 + b.0, a.1 + b.1)
+    }
+    fn identity() -> Self::M {
+        (0, 0)
+    }
+}
+
+
 #[cfg(test)]
 mod test {
     use super::*;
 
     #[test]
     fn test_arq_tree() {
-        let mut arq = ArqTree::new(10);
+        let mut arq = ArqTree::<AssignAdd>::new(vec![(0, 1); 10]);
+
+        assert_eq!(arq.query(0, 9), (0, 10));
 
-        arq.modify(1, 3, 10);
-        arq.modify(3, 5, 1);
+        arq.modify(1, 3, &10);
+        arq.modify(3, 5, &1);
 
-        assert_eq!(arq.query(0, 9), 23);
+        assert_eq!(arq.query(0, 9), (23, 10));
     }
 }