contest_algorithms/range_query/
static_arq.rs

1//! Associative Range Query Tree
2use super::ArqSpec;
3
4/// Colloquially known as a "segtree" in the sport programming literature, it
5/// represents a sequence of elements a_i (0 <= i < size) from a monoid (S, +)
6/// on which we want to support fast range operations:
7///
8/// - update(l, r, f) replaces a_i (l <= i <= r) by f(a_i) for an endomorphism f
9/// - query(l, r) returns the aggregate a_l + a_{l+1} + ... + a_r
10///
11/// This compact representation is based on a [blog post by Al.Cash]
12/// (http://codeforces.com/blog/entry/18051). All nodes have 0 or 2 children.
13/// Hence, trees whose size is not a power of two will have multiple roots.
14///
15/// Future work: ArqTree would lend itself naturally to Rust's ownership system.
16/// Initially, we should only have access to the root nodes:
17///            if size is a power of two, there is a unique root at index 1.
18/// arq.push(i) locks i and acquires access to its children.
19/// arq.pull(i) is called when the lock on i is released.
20pub struct StaticArq<T: ArqSpec> {
21    val: Vec<T::S>,
22    app: Vec<Option<T::F>>,
23}
24
25impl<T: ArqSpec> StaticArq<T> {
26    /// Initializes a static balanced binary tree on top of the given sequence.
27    pub fn new(init_val: &[T::S]) -> Self {
28        let size = init_val.len();
29        let mut val = vec![T::identity(); size];
30        val.extend_from_slice(init_val);
31        let app = vec![None; size];
32
33        let mut arq = Self { val, app };
34        for p in (0..size).rev() {
35            arq.pull(p);
36        }
37        arq
38    }
39
40    fn apply(&mut self, p: usize, f: &T::F, s: i64) {
41        self.val[p] = T::apply(f, &self.val[p], s);
42        if let Some(lazy) = self.app.get_mut(p) {
43            let h = match *lazy {
44                Some(ref g) => T::compose(f, g),
45                None => f.clone(),
46            };
47            *lazy = Some(h);
48        }
49    }
50
51    fn push(&mut self, p: usize) {
52        if let Some(ref f) = self.app[p].take() {
53            let s = ((self.app.len() + p - 1) / p / 2).next_power_of_two() as i64;
54            self.apply(p << 1, f, s);
55            self.apply(p << 1 | 1, f, s);
56        }
57    }
58
59    fn pull(&mut self, p: usize) {
60        self.val[p] = T::op(&self.val[p << 1], &self.val[p << 1 | 1]);
61    }
62
63    fn push_to(&mut self, p: usize) {
64        let one_plus_floor_log_p = (p + 1).next_power_of_two().trailing_zeros();
65        for i in (1..one_plus_floor_log_p).rev() {
66            self.push(p >> i);
67        }
68    }
69
70    fn pull_from(&mut self, mut p: usize) {
71        while p > 1 {
72            p >>= 1;
73            self.pull(p);
74        }
75    }
76
77    /// Applies the endomorphism f to all entries from l to r, inclusive.
78    /// If l == r, the updates are eager. Otherwise, they are lazy.
79    ///
80    /// # Panics
81    ///
82    /// Panics if r >= size. Note that l > r is valid, meaning an empty range.
83    pub fn update(&mut self, mut l: usize, mut r: usize, f: &T::F) {
84        l += self.app.len();
85        r += self.app.len();
86        if l < r {
87            self.push_to(l);
88        }
89        self.push_to(r);
90        let (mut l0, mut r0, mut s) = (1, 1, 1);
91        while l <= r {
92            if l & 1 == 1 {
93                self.apply(l, f, s);
94                l0 = l0.max(l);
95                l += 1;
96            }
97            if r & 1 == 0 {
98                self.apply(r, f, s);
99                r0 = r0.max(r);
100                r -= 1;
101            }
102            l >>= 1;
103            r >>= 1;
104            s <<= 1;
105        }
106        self.pull_from(l0);
107        self.pull_from(r0);
108    }
109
110    /// Returns the aggregate range query on all entries from l to r, inclusive.
111    ///
112    /// # Panics
113    ///
114    /// Panics if r >= size. Note that l > r is valid, meaning an empty range.
115    pub fn query(&mut self, mut l: usize, mut r: usize) -> T::S {
116        l += self.app.len();
117        r += self.app.len();
118        if l < r {
119            self.push_to(l);
120        }
121        self.push_to(r);
122        let (mut l_agg, mut r_agg) = (T::identity(), T::identity());
123        while l <= r {
124            if l & 1 == 1 {
125                l_agg = T::op(&l_agg, &self.val[l]);
126                l += 1;
127            }
128            if r & 1 == 0 {
129                r_agg = T::op(&self.val[r], &r_agg);
130                r -= 1;
131            }
132            l >>= 1;
133            r >>= 1;
134        }
135        T::op(&l_agg, &r_agg)
136    }
137}
138
139/// An example of binary search to find the first position whose element is negative.
140/// In this case, we use RMQ to locate the leftmost negative element.
141/// To ensure the existence of a valid root note (i == 1) from which to descend,
142/// the tree's size must be a power of two.
143pub fn first_negative(arq: &mut StaticArq<super::specs::AssignMin>) -> Option<usize> {
144    assert!(arq.app.len().is_power_of_two());
145    let mut p = 1;
146    if arq.val[p] >= 0 {
147        None
148    } else {
149        while p < arq.app.len() {
150            arq.push(p);
151            p <<= 1;
152            if arq.val[p] >= 0 {
153                p |= 1;
154            }
155        }
156        Some(p - arq.app.len())
157    }
158}
contest_algorithms/range_query/static_arq.rs

contest_algorithms/range_query/
static_arq.rs
