Struct WaveletMatrix

Source

pub struct WaveletMatrix<T> {
    len: usize,
    bit_length: usize,
    zeros: Vec<usize>,
    ones_prefix: Vec<usize>,
    bit_vector: BitVector,
    compress: VecCompress<T>,
}

Fields§

§len: usize§bit_length: usize§zeros: Vec<usize>§ones_prefix: Vec<usize>§bit_vector: BitVector§compress: VecCompress<T>

Implementations§

Source §

impl<T> WaveletMatrix<T>
where T: Ord + Clone,

Source

pub fn new(v: Vec<T>) -> Self

Examples found in repository ?

crates/library_checker/src/data_structure/range_kth_smallest.rs (line 10)

6pub fn range_kth_smallest(reader: impl Read, mut writer: impl Write) {
7    let s = read_all_unchecked(reader);
8    let mut scanner = Scanner::new(&s);
9    scan!(scanner, n, q, a: [usize; n]);
10    let wm = WaveletMatrix::new(a);
11    for (l, r, k) in scanner.iter::<(usize, usize, usize)>().take(q) {
12        writeln!(writer, "{}", wm.quantile(l..r, k)).ok();
13    }
14}

More examples

Hide additional examples

crates/library_checker/src/data_structure/static_range_frequency.rs (line 10)

6pub fn static_range_frequency(reader: impl Read, mut writer: impl Write) {
7    let s = read_all_unchecked(reader);
8    let mut scanner = Scanner::new(&s);
9    scan!(scanner, n, q, a: [usize; n]);
10    let wm = WaveletMatrix::new(a);
11    for _ in 0..q {
12        scan!(scanner, l, r, x: usize);
13        let ans = wm.rank(x, l..r);
14        iter_print!(writer, ans);
15    }
16}

crates/library_checker/src/data_structure/static_range_sum_with_upper_bound.rs (line 11)

6pub fn static_range_sum_with_upper_bound(reader: impl Read, mut writer: impl Write) {
7    let s = read_all_unchecked(reader);
8    let mut scanner = Scanner::new(&s);
9    scan!(scanner, n, q, a: [i64; n]);
10    let weights = a.clone();
11    let wm = WaveletMatrix::new(a);
12    let fold = wm.build_fold::<AdditiveOperation<i64>>(&weights);
13    for _ in 0..q {
14        scan!(scanner, l, r, x: i64);
15        let (count, sum) = fold.fold_lessthan_with_count(x + 1, l..r);
16        writeln!(writer, "{} {}", count, sum).ok();
17    }
18}

crates/competitive/src/data_structure/wavelet_matrix.rs (line 76)

72    pub fn new_with_init<F>(v: Vec<T>, mut f: F) -> Self
73    where
74        F: FnMut(usize, usize, T),
75    {
76        let this = Self::new(v.clone());
77        let indices: Vec<usize> = v
78            .iter()
79            .map(|value| this.compress.index_exact(value).unwrap())
80            .collect();
81        for (mut k, value) in v.into_iter().enumerate() {
82            let idx = indices[k];
83            for d in (0..this.bit_length).rev() {
84                let level = this.level(d);
85                if ((idx >> d) & 1) != 0 {
86                    k = this.zeros[level] + this.rank1(level, k);
87                } else {
88                    k = this.rank0(level, k);
89                }
90                f(d, k, value.clone());
91            }
92        }
93        this
94    }

Source

pub fn new_with_init<F>(v: Vec<T>, f: F) -> Self
where F: FnMut(usize, usize, T),

Source

fn level(&self, d: usize) -> usize

Examples found in repository ?

crates/competitive/src/data_structure/wavelet_matrix.rs (line 84)

72    pub fn new_with_init<F>(v: Vec<T>, mut f: F) -> Self
73    where
74        F: FnMut(usize, usize, T),
75    {
76        let this = Self::new(v.clone());
77        let indices: Vec<usize> = v
78            .iter()
79            .map(|value| this.compress.index_exact(value).unwrap())
80            .collect();
81        for (mut k, value) in v.into_iter().enumerate() {
82            let idx = indices[k];
83            for d in (0..this.bit_length).rev() {
84                let level = this.level(d);
85                if ((idx >> d) & 1) != 0 {
86                    k = this.zeros[level] + this.rank1(level, k);
87                } else {
88                    k = this.rank0(level, k);
89                }
90                f(d, k, value.clone());
91            }
92        }
93        this
94    }
95
96    fn level(&self, d: usize) -> usize {
97        self.bit_length - 1 - d
98    }
99
100    fn rank1(&self, level: usize, k: usize) -> usize {
101        let offset = level * self.len;
102        self.bit_vector.rank1(offset + k) - self.ones_prefix[level]
103    }
104
105    fn rank0(&self, level: usize, k: usize) -> usize {
106        k - self.rank1(level, k)
107    }
108
109    fn rank_by_index(&self, idx: usize, mut range: Range<usize>) -> usize {
110        for d in (0..self.bit_length).rev() {
111            let level = self.level(d);
112            if ((idx >> d) & 1) != 0 {
113                range.start = self.zeros[level] + self.rank1(level, range.start);
114                range.end = self.zeros[level] + self.rank1(level, range.end);
115            } else {
116                range.start = self.rank0(level, range.start);
117                range.end = self.rank0(level, range.end);
118            }
119        }
120        range.end - range.start
121    }
122
123    /// get k-th value
124    pub fn access(&self, mut k: usize) -> T {
125        let mut idx = 0;
126        for d in (0..self.bit_length).rev() {
127            let level = self.level(d);
128            if self.bit_vector.access(level * self.len + k) {
129                idx |= 1 << d;
130                k = self.zeros[level] + self.rank1(level, k);
131            } else {
132                k = self.rank0(level, k);
133            }
134        }
135        self.compress.values()[idx].clone()
136    }
137
138    /// the number of val in range
139    pub fn rank(&self, val: T, range: Range<usize>) -> usize {
140        match self.compress.index_exact(&val) {
141            Some(idx) => self.rank_by_index(idx, range),
142            None => 0,
143        }
144    }
145
146    /// index of k-th val
147    pub fn select(&self, val: T, k: usize) -> Option<usize> {
148        let idx = self.compress.index_exact(&val)?;
149        if self.rank_by_index(idx, 0..self.len) <= k {
150            return None;
151        }
152        let mut i = 0;
153        for d in (0..self.bit_length).rev() {
154            let level = self.level(d);
155            if ((idx >> d) & 1) != 0 {
156                i = self.zeros[level] + self.rank1(level, i);
157            } else {
158                i = self.rank0(level, i);
159            }
160        }
161        i += k;
162        for level in (0..self.bit_length).rev() {
163            let offset = level * self.len;
164            if i >= self.zeros[level] {
165                let global_k = self.ones_prefix[level] + (i - self.zeros[level]);
166                let pos = self.bit_vector.select1(global_k).unwrap();
167                i = pos - offset;
168            } else {
169                let zeros_before = offset - self.ones_prefix[level];
170                let global_k = zeros_before + i;
171                let pos = self.bit_vector.select0(global_k).unwrap();
172                i = pos - offset;
173            }
174        }
175        Some(i)
176    }
177
178    /// get k-th smallest value in range
179    pub fn quantile(&self, mut range: Range<usize>, mut k: usize) -> T {
180        let mut idx = 0;
181        for d in (0..self.bit_length).rev() {
182            let level = self.level(d);
183            let z = self.rank0(level, range.end) - self.rank0(level, range.start);
184            if z <= k {
185                k -= z;
186                idx |= 1 << d;
187                range.start = self.zeros[level] + self.rank1(level, range.start);
188                range.end = self.zeros[level] + self.rank1(level, range.end);
189            } else {
190                range.start = self.rank0(level, range.start);
191                range.end = self.rank0(level, range.end);
192            }
193        }
194        self.compress.values()[idx].clone()
195    }
196
197    /// get k-th smallest value out of range
198    pub fn quantile_outer(&self, mut range: Range<usize>, mut k: usize) -> T {
199        let mut idx = 0;
200        let mut orange = 0..self.len;
201        for d in (0..self.bit_length).rev() {
202            let level = self.level(d);
203            let z = self.rank0(level, orange.end) - self.rank0(level, orange.start)
204                + self.rank0(level, range.start)
205                - self.rank0(level, range.end);
206            if z <= k {
207                k -= z;
208                idx |= 1 << d;
209                range.start = self.zeros[level] + self.rank1(level, range.start);
210                range.end = self.zeros[level] + self.rank1(level, range.end);
211                orange.start = self.zeros[level] + self.rank1(level, orange.start);
212                orange.end = self.zeros[level] + self.rank1(level, orange.end);
213            } else {
214                range.start = self.rank0(level, range.start);
215                range.end = self.rank0(level, range.end);
216                orange.start = self.rank0(level, orange.start);
217                orange.end = self.rank0(level, orange.end);
218            }
219        }
220        self.compress.values()[idx].clone()
221    }
222
223    /// the number of value less than val in range
224    pub fn rank_lessthan(&self, val: T, mut range: Range<usize>) -> usize {
225        let idx = self.compress.index_lower_bound(&val);
226        let mut res = 0;
227        for d in (0..self.bit_length).rev() {
228            let level = self.level(d);
229            if ((idx >> d) & 1) != 0 {
230                res += self.rank0(level, range.end) - self.rank0(level, range.start);
231                range.start = self.zeros[level] + self.rank1(level, range.start);
232                range.end = self.zeros[level] + self.rank1(level, range.end);
233            } else {
234                range.start = self.rank0(level, range.start);
235                range.end = self.rank0(level, range.end);
236            }
237        }
238        res
239    }
240
241    /// the number of valrange in range
242    pub fn rank_range(&self, valrange: Range<T>, range: Range<usize>) -> usize {
243        self.rank_lessthan(valrange.end, range.clone()) - self.rank_lessthan(valrange.start, range)
244    }
245
246    pub fn query_less_than<F>(&self, val: T, mut range: Range<usize>, mut f: F)
247    where
248        F: FnMut(usize, Range<usize>),
249    {
250        let idx = self.compress.index_lower_bound(&val);
251        for d in (0..self.bit_length).rev() {
252            let level = self.level(d);
253            if ((idx >> d) & 1) != 0 {
254                f(
255                    d,
256                    self.rank0(level, range.start)..self.rank0(level, range.end),
257                );
258                range.start = self.zeros[level] + self.rank1(level, range.start);
259                range.end = self.zeros[level] + self.rank1(level, range.end);
260            } else {
261                range.start = self.rank0(level, range.start);
262                range.end = self.rank0(level, range.end);
263            }
264        }
265    }
266
267    pub fn build_fold<M>(&self, weights: &[M::T]) -> WaveletMatrixFold<'_, T, M>
268    where
269        M: AbelianGroup,
270    {
271        let len = self.len;
272        assert_eq!(weights.len(), len);
273        let mut prefix = Vec::with_capacity((self.bit_length + 1) * (len + 1));
274        let mut current: Vec<M::T> = weights.to_vec();
275        for level in 0..self.bit_length {
276            let offset = level * len;
277            let zeros = self.zeros[level];
278            let mut next: Vec<MaybeUninit<M::T>> = Vec::with_capacity(len);
279            next.resize_with(len, MaybeUninit::uninit);
280            let mut zero_pos = 0;
281            let mut one_pos = zeros;
282            let mut acc = M::unit();
283            prefix.push(acc.clone());
284            for (i, w) in current.into_iter().enumerate() {
285                acc = M::operate(&acc, &w);
286                prefix.push(acc.clone());
287                if self.bit_vector.access(offset + i) {
288                    next[one_pos].write(w);
289                    one_pos += 1;
290                } else {
291                    next[zero_pos].write(w);
292                    zero_pos += 1;
293                }
294            }
295            debug_assert_eq!(zero_pos, zeros);
296            debug_assert_eq!(one_pos, len);
297            let next = unsafe {
298                let mut next = mem::ManuallyDrop::new(next);
299                let ptr = next.as_mut_ptr() as *mut M::T;
300                let len = next.len();
301                let cap = next.capacity();
302                Vec::from_raw_parts(ptr, len, cap)
303            };
304            current = next;
305        }
306        let mut acc = M::unit();
307        prefix.push(acc.clone());
308        for w in current.into_iter() {
309            acc = M::operate(&acc, &w);
310            prefix.push(acc.clone());
311        }
312        WaveletMatrixFold {
313            wavelet_matrix: self,
314            prefix,
315        }
316    }
317}
318
319#[derive(Debug, Clone)]
320pub struct WaveletMatrixFold<'a, T, M>
321where
322    T: Ord + Clone,
323    M: AbelianGroup,
324{
325    wavelet_matrix: &'a WaveletMatrix<T>,
326    prefix: Vec<M::T>,
327}
328
329impl<'a, T, M> WaveletMatrixFold<'a, T, M>
330where
331    T: Ord + Clone,
332    M: AbelianGroup,
333{
334    #[inline]
335    fn range_sum(&self, level: usize, range: Range<usize>) -> M::T {
336        let offset = level * (self.wavelet_matrix.len + 1);
337        unsafe {
338            M::rinv_operate(
339                self.prefix.get_unchecked(offset + range.end),
340                self.prefix.get_unchecked(offset + range.start),
341            )
342        }
343    }
344
345    pub fn fold_lessthan(&self, val: T, range: Range<usize>) -> M::T {
346        self.fold_lessthan_with_count(val, range).1
347    }
348
349    pub fn fold_lessthan_with_count(&self, val: T, mut range: Range<usize>) -> (usize, M::T) {
350        debug_assert!(range.end <= self.wavelet_matrix.len);
351        let idx = self.wavelet_matrix.compress.index_lower_bound(&val);
352        let mut count = 0;
353        let mut sum = M::unit();
354        for d in (0..self.wavelet_matrix.bit_length).rev() {
355            let level = self.wavelet_matrix.level(d);
356            let start0 = self.wavelet_matrix.rank0(level, range.start);
357            let end0 = self.wavelet_matrix.rank0(level, range.end);
358            if ((idx >> d) & 1) != 0 {
359                count += end0 - start0;
360                sum = M::operate(&sum, &self.range_sum(level + 1, start0..end0));
361                range.start = self.wavelet_matrix.zeros[level] + (range.start - start0);
362                range.end = self.wavelet_matrix.zeros[level] + (range.end - end0);
363            } else {
364                range.start = start0;
365                range.end = end0;
366            }
367        }
368        (count, sum)
369    }

Source

fn rank1(&self, level: usize, k: usize) -> usize

Examples found in repository ?

crates/competitive/src/data_structure/wavelet_matrix.rs (line 86)

72    pub fn new_with_init<F>(v: Vec<T>, mut f: F) -> Self
73    where
74        F: FnMut(usize, usize, T),
75    {
76        let this = Self::new(v.clone());
77        let indices: Vec<usize> = v
78            .iter()
79            .map(|value| this.compress.index_exact(value).unwrap())
80            .collect();
81        for (mut k, value) in v.into_iter().enumerate() {
82            let idx = indices[k];
83            for d in (0..this.bit_length).rev() {
84                let level = this.level(d);
85                if ((idx >> d) & 1) != 0 {
86                    k = this.zeros[level] + this.rank1(level, k);
87                } else {
88                    k = this.rank0(level, k);
89                }
90                f(d, k, value.clone());
91            }
92        }
93        this
94    }
95
96    fn level(&self, d: usize) -> usize {
97        self.bit_length - 1 - d
98    }
99
100    fn rank1(&self, level: usize, k: usize) -> usize {
101        let offset = level * self.len;
102        self.bit_vector.rank1(offset + k) - self.ones_prefix[level]
103    }
104
105    fn rank0(&self, level: usize, k: usize) -> usize {
106        k - self.rank1(level, k)
107    }
108
109    fn rank_by_index(&self, idx: usize, mut range: Range<usize>) -> usize {
110        for d in (0..self.bit_length).rev() {
111            let level = self.level(d);
112            if ((idx >> d) & 1) != 0 {
113                range.start = self.zeros[level] + self.rank1(level, range.start);
114                range.end = self.zeros[level] + self.rank1(level, range.end);
115            } else {
116                range.start = self.rank0(level, range.start);
117                range.end = self.rank0(level, range.end);
118            }
119        }
120        range.end - range.start
121    }
122
123    /// get k-th value
124    pub fn access(&self, mut k: usize) -> T {
125        let mut idx = 0;
126        for d in (0..self.bit_length).rev() {
127            let level = self.level(d);
128            if self.bit_vector.access(level * self.len + k) {
129                idx |= 1 << d;
130                k = self.zeros[level] + self.rank1(level, k);
131            } else {
132                k = self.rank0(level, k);
133            }
134        }
135        self.compress.values()[idx].clone()
136    }
137
138    /// the number of val in range
139    pub fn rank(&self, val: T, range: Range<usize>) -> usize {
140        match self.compress.index_exact(&val) {
141            Some(idx) => self.rank_by_index(idx, range),
142            None => 0,
143        }
144    }
145
146    /// index of k-th val
147    pub fn select(&self, val: T, k: usize) -> Option<usize> {
148        let idx = self.compress.index_exact(&val)?;
149        if self.rank_by_index(idx, 0..self.len) <= k {
150            return None;
151        }
152        let mut i = 0;
153        for d in (0..self.bit_length).rev() {
154            let level = self.level(d);
155            if ((idx >> d) & 1) != 0 {
156                i = self.zeros[level] + self.rank1(level, i);
157            } else {
158                i = self.rank0(level, i);
159            }
160        }
161        i += k;
162        for level in (0..self.bit_length).rev() {
163            let offset = level * self.len;
164            if i >= self.zeros[level] {
165                let global_k = self.ones_prefix[level] + (i - self.zeros[level]);
166                let pos = self.bit_vector.select1(global_k).unwrap();
167                i = pos - offset;
168            } else {
169                let zeros_before = offset - self.ones_prefix[level];
170                let global_k = zeros_before + i;
171                let pos = self.bit_vector.select0(global_k).unwrap();
172                i = pos - offset;
173            }
174        }
175        Some(i)
176    }
177
178    /// get k-th smallest value in range
179    pub fn quantile(&self, mut range: Range<usize>, mut k: usize) -> T {
180        let mut idx = 0;
181        for d in (0..self.bit_length).rev() {
182            let level = self.level(d);
183            let z = self.rank0(level, range.end) - self.rank0(level, range.start);
184            if z <= k {
185                k -= z;
186                idx |= 1 << d;
187                range.start = self.zeros[level] + self.rank1(level, range.start);
188                range.end = self.zeros[level] + self.rank1(level, range.end);
189            } else {
190                range.start = self.rank0(level, range.start);
191                range.end = self.rank0(level, range.end);
192            }
193        }
194        self.compress.values()[idx].clone()
195    }
196
197    /// get k-th smallest value out of range
198    pub fn quantile_outer(&self, mut range: Range<usize>, mut k: usize) -> T {
199        let mut idx = 0;
200        let mut orange = 0..self.len;
201        for d in (0..self.bit_length).rev() {
202            let level = self.level(d);
203            let z = self.rank0(level, orange.end) - self.rank0(level, orange.start)
204                + self.rank0(level, range.start)
205                - self.rank0(level, range.end);
206            if z <= k {
207                k -= z;
208                idx |= 1 << d;
209                range.start = self.zeros[level] + self.rank1(level, range.start);
210                range.end = self.zeros[level] + self.rank1(level, range.end);
211                orange.start = self.zeros[level] + self.rank1(level, orange.start);
212                orange.end = self.zeros[level] + self.rank1(level, orange.end);
213            } else {
214                range.start = self.rank0(level, range.start);
215                range.end = self.rank0(level, range.end);
216                orange.start = self.rank0(level, orange.start);
217                orange.end = self.rank0(level, orange.end);
218            }
219        }
220        self.compress.values()[idx].clone()
221    }
222
223    /// the number of value less than val in range
224    pub fn rank_lessthan(&self, val: T, mut range: Range<usize>) -> usize {
225        let idx = self.compress.index_lower_bound(&val);
226        let mut res = 0;
227        for d in (0..self.bit_length).rev() {
228            let level = self.level(d);
229            if ((idx >> d) & 1) != 0 {
230                res += self.rank0(level, range.end) - self.rank0(level, range.start);
231                range.start = self.zeros[level] + self.rank1(level, range.start);
232                range.end = self.zeros[level] + self.rank1(level, range.end);
233            } else {
234                range.start = self.rank0(level, range.start);
235                range.end = self.rank0(level, range.end);
236            }
237        }
238        res
239    }
240
241    /// the number of valrange in range
242    pub fn rank_range(&self, valrange: Range<T>, range: Range<usize>) -> usize {
243        self.rank_lessthan(valrange.end, range.clone()) - self.rank_lessthan(valrange.start, range)
244    }
245
246    pub fn query_less_than<F>(&self, val: T, mut range: Range<usize>, mut f: F)
247    where
248        F: FnMut(usize, Range<usize>),
249    {
250        let idx = self.compress.index_lower_bound(&val);
251        for d in (0..self.bit_length).rev() {
252            let level = self.level(d);
253            if ((idx >> d) & 1) != 0 {
254                f(
255                    d,
256                    self.rank0(level, range.start)..self.rank0(level, range.end),
257                );
258                range.start = self.zeros[level] + self.rank1(level, range.start);
259                range.end = self.zeros[level] + self.rank1(level, range.end);
260            } else {
261                range.start = self.rank0(level, range.start);
262                range.end = self.rank0(level, range.end);
263            }
264        }
265    }

Source

fn rank0(&self, level: usize, k: usize) -> usize

Examples found in repository ?

crates/competitive/src/data_structure/wavelet_matrix.rs (line 88)

72    pub fn new_with_init<F>(v: Vec<T>, mut f: F) -> Self
73    where
74        F: FnMut(usize, usize, T),
75    {
76        let this = Self::new(v.clone());
77        let indices: Vec<usize> = v
78            .iter()
79            .map(|value| this.compress.index_exact(value).unwrap())
80            .collect();
81        for (mut k, value) in v.into_iter().enumerate() {
82            let idx = indices[k];
83            for d in (0..this.bit_length).rev() {
84                let level = this.level(d);
85                if ((idx >> d) & 1) != 0 {
86                    k = this.zeros[level] + this.rank1(level, k);
87                } else {
88                    k = this.rank0(level, k);
89                }
90                f(d, k, value.clone());
91            }
92        }
93        this
94    }
95
96    fn level(&self, d: usize) -> usize {
97        self.bit_length - 1 - d
98    }
99
100    fn rank1(&self, level: usize, k: usize) -> usize {
101        let offset = level * self.len;
102        self.bit_vector.rank1(offset + k) - self.ones_prefix[level]
103    }
104
105    fn rank0(&self, level: usize, k: usize) -> usize {
106        k - self.rank1(level, k)
107    }
108
109    fn rank_by_index(&self, idx: usize, mut range: Range<usize>) -> usize {
110        for d in (0..self.bit_length).rev() {
111            let level = self.level(d);
112            if ((idx >> d) & 1) != 0 {
113                range.start = self.zeros[level] + self.rank1(level, range.start);
114                range.end = self.zeros[level] + self.rank1(level, range.end);
115            } else {
116                range.start = self.rank0(level, range.start);
117                range.end = self.rank0(level, range.end);
118            }
119        }
120        range.end - range.start
121    }
122
123    /// get k-th value
124    pub fn access(&self, mut k: usize) -> T {
125        let mut idx = 0;
126        for d in (0..self.bit_length).rev() {
127            let level = self.level(d);
128            if self.bit_vector.access(level * self.len + k) {
129                idx |= 1 << d;
130                k = self.zeros[level] + self.rank1(level, k);
131            } else {
132                k = self.rank0(level, k);
133            }
134        }
135        self.compress.values()[idx].clone()
136    }
137
138    /// the number of val in range
139    pub fn rank(&self, val: T, range: Range<usize>) -> usize {
140        match self.compress.index_exact(&val) {
141            Some(idx) => self.rank_by_index(idx, range),
142            None => 0,
143        }
144    }
145
146    /// index of k-th val
147    pub fn select(&self, val: T, k: usize) -> Option<usize> {
148        let idx = self.compress.index_exact(&val)?;
149        if self.rank_by_index(idx, 0..self.len) <= k {
150            return None;
151        }
152        let mut i = 0;
153        for d in (0..self.bit_length).rev() {
154            let level = self.level(d);
155            if ((idx >> d) & 1) != 0 {
156                i = self.zeros[level] + self.rank1(level, i);
157            } else {
158                i = self.rank0(level, i);
159            }
160        }
161        i += k;
162        for level in (0..self.bit_length).rev() {
163            let offset = level * self.len;
164            if i >= self.zeros[level] {
165                let global_k = self.ones_prefix[level] + (i - self.zeros[level]);
166                let pos = self.bit_vector.select1(global_k).unwrap();
167                i = pos - offset;
168            } else {
169                let zeros_before = offset - self.ones_prefix[level];
170                let global_k = zeros_before + i;
171                let pos = self.bit_vector.select0(global_k).unwrap();
172                i = pos - offset;
173            }
174        }
175        Some(i)
176    }
177
178    /// get k-th smallest value in range
179    pub fn quantile(&self, mut range: Range<usize>, mut k: usize) -> T {
180        let mut idx = 0;
181        for d in (0..self.bit_length).rev() {
182            let level = self.level(d);
183            let z = self.rank0(level, range.end) - self.rank0(level, range.start);
184            if z <= k {
185                k -= z;
186                idx |= 1 << d;
187                range.start = self.zeros[level] + self.rank1(level, range.start);
188                range.end = self.zeros[level] + self.rank1(level, range.end);
189            } else {
190                range.start = self.rank0(level, range.start);
191                range.end = self.rank0(level, range.end);
192            }
193        }
194        self.compress.values()[idx].clone()
195    }
196
197    /// get k-th smallest value out of range
198    pub fn quantile_outer(&self, mut range: Range<usize>, mut k: usize) -> T {
199        let mut idx = 0;
200        let mut orange = 0..self.len;
201        for d in (0..self.bit_length).rev() {
202            let level = self.level(d);
203            let z = self.rank0(level, orange.end) - self.rank0(level, orange.start)
204                + self.rank0(level, range.start)
205                - self.rank0(level, range.end);
206            if z <= k {
207                k -= z;
208                idx |= 1 << d;
209                range.start = self.zeros[level] + self.rank1(level, range.start);
210                range.end = self.zeros[level] + self.rank1(level, range.end);
211                orange.start = self.zeros[level] + self.rank1(level, orange.start);
212                orange.end = self.zeros[level] + self.rank1(level, orange.end);
213            } else {
214                range.start = self.rank0(level, range.start);
215                range.end = self.rank0(level, range.end);
216                orange.start = self.rank0(level, orange.start);
217                orange.end = self.rank0(level, orange.end);
218            }
219        }
220        self.compress.values()[idx].clone()
221    }
222
223    /// the number of value less than val in range
224    pub fn rank_lessthan(&self, val: T, mut range: Range<usize>) -> usize {
225        let idx = self.compress.index_lower_bound(&val);
226        let mut res = 0;
227        for d in (0..self.bit_length).rev() {
228            let level = self.level(d);
229            if ((idx >> d) & 1) != 0 {
230                res += self.rank0(level, range.end) - self.rank0(level, range.start);
231                range.start = self.zeros[level] + self.rank1(level, range.start);
232                range.end = self.zeros[level] + self.rank1(level, range.end);
233            } else {
234                range.start = self.rank0(level, range.start);
235                range.end = self.rank0(level, range.end);
236            }
237        }
238        res
239    }
240
241    /// the number of valrange in range
242    pub fn rank_range(&self, valrange: Range<T>, range: Range<usize>) -> usize {
243        self.rank_lessthan(valrange.end, range.clone()) - self.rank_lessthan(valrange.start, range)
244    }
245
246    pub fn query_less_than<F>(&self, val: T, mut range: Range<usize>, mut f: F)
247    where
248        F: FnMut(usize, Range<usize>),
249    {
250        let idx = self.compress.index_lower_bound(&val);
251        for d in (0..self.bit_length).rev() {
252            let level = self.level(d);
253            if ((idx >> d) & 1) != 0 {
254                f(
255                    d,
256                    self.rank0(level, range.start)..self.rank0(level, range.end),
257                );
258                range.start = self.zeros[level] + self.rank1(level, range.start);
259                range.end = self.zeros[level] + self.rank1(level, range.end);
260            } else {
261                range.start = self.rank0(level, range.start);
262                range.end = self.rank0(level, range.end);
263            }
264        }
265    }
266
267    pub fn build_fold<M>(&self, weights: &[M::T]) -> WaveletMatrixFold<'_, T, M>
268    where
269        M: AbelianGroup,
270    {
271        let len = self.len;
272        assert_eq!(weights.len(), len);
273        let mut prefix = Vec::with_capacity((self.bit_length + 1) * (len + 1));
274        let mut current: Vec<M::T> = weights.to_vec();
275        for level in 0..self.bit_length {
276            let offset = level * len;
277            let zeros = self.zeros[level];
278            let mut next: Vec<MaybeUninit<M::T>> = Vec::with_capacity(len);
279            next.resize_with(len, MaybeUninit::uninit);
280            let mut zero_pos = 0;
281            let mut one_pos = zeros;
282            let mut acc = M::unit();
283            prefix.push(acc.clone());
284            for (i, w) in current.into_iter().enumerate() {
285                acc = M::operate(&acc, &w);
286                prefix.push(acc.clone());
287                if self.bit_vector.access(offset + i) {
288                    next[one_pos].write(w);
289                    one_pos += 1;
290                } else {
291                    next[zero_pos].write(w);
292                    zero_pos += 1;
293                }
294            }
295            debug_assert_eq!(zero_pos, zeros);
296            debug_assert_eq!(one_pos, len);
297            let next = unsafe {
298                let mut next = mem::ManuallyDrop::new(next);
299                let ptr = next.as_mut_ptr() as *mut M::T;
300                let len = next.len();
301                let cap = next.capacity();
302                Vec::from_raw_parts(ptr, len, cap)
303            };
304            current = next;
305        }
306        let mut acc = M::unit();
307        prefix.push(acc.clone());
308        for w in current.into_iter() {
309            acc = M::operate(&acc, &w);
310            prefix.push(acc.clone());
311        }
312        WaveletMatrixFold {
313            wavelet_matrix: self,
314            prefix,
315        }
316    }
317}
318
319#[derive(Debug, Clone)]
320pub struct WaveletMatrixFold<'a, T, M>
321where
322    T: Ord + Clone,
323    M: AbelianGroup,
324{
325    wavelet_matrix: &'a WaveletMatrix<T>,
326    prefix: Vec<M::T>,
327}
328
329impl<'a, T, M> WaveletMatrixFold<'a, T, M>
330where
331    T: Ord + Clone,
332    M: AbelianGroup,
333{
334    #[inline]
335    fn range_sum(&self, level: usize, range: Range<usize>) -> M::T {
336        let offset = level * (self.wavelet_matrix.len + 1);
337        unsafe {
338            M::rinv_operate(
339                self.prefix.get_unchecked(offset + range.end),
340                self.prefix.get_unchecked(offset + range.start),
341            )
342        }
343    }
344
345    pub fn fold_lessthan(&self, val: T, range: Range<usize>) -> M::T {
346        self.fold_lessthan_with_count(val, range).1
347    }
348
349    pub fn fold_lessthan_with_count(&self, val: T, mut range: Range<usize>) -> (usize, M::T) {
350        debug_assert!(range.end <= self.wavelet_matrix.len);
351        let idx = self.wavelet_matrix.compress.index_lower_bound(&val);
352        let mut count = 0;
353        let mut sum = M::unit();
354        for d in (0..self.wavelet_matrix.bit_length).rev() {
355            let level = self.wavelet_matrix.level(d);
356            let start0 = self.wavelet_matrix.rank0(level, range.start);
357            let end0 = self.wavelet_matrix.rank0(level, range.end);
358            if ((idx >> d) & 1) != 0 {
359                count += end0 - start0;
360                sum = M::operate(&sum, &self.range_sum(level + 1, start0..end0));
361                range.start = self.wavelet_matrix.zeros[level] + (range.start - start0);
362                range.end = self.wavelet_matrix.zeros[level] + (range.end - end0);
363            } else {
364                range.start = start0;
365                range.end = end0;
366            }
367        }
368        (count, sum)
369    }

Source

fn rank_by_index(&self, idx: usize, range: Range<usize>) -> usize

Examples found in repository ?

crates/competitive/src/data_structure/wavelet_matrix.rs (line 141)

139    pub fn rank(&self, val: T, range: Range<usize>) -> usize {
140        match self.compress.index_exact(&val) {
141            Some(idx) => self.rank_by_index(idx, range),
142            None => 0,
143        }
144    }
145
146    /// index of k-th val
147    pub fn select(&self, val: T, k: usize) -> Option<usize> {
148        let idx = self.compress.index_exact(&val)?;
149        if self.rank_by_index(idx, 0..self.len) <= k {
150            return None;
151        }
152        let mut i = 0;
153        for d in (0..self.bit_length).rev() {
154            let level = self.level(d);
155            if ((idx >> d) & 1) != 0 {
156                i = self.zeros[level] + self.rank1(level, i);
157            } else {
158                i = self.rank0(level, i);
159            }
160        }
161        i += k;
162        for level in (0..self.bit_length).rev() {
163            let offset = level * self.len;
164            if i >= self.zeros[level] {
165                let global_k = self.ones_prefix[level] + (i - self.zeros[level]);
166                let pos = self.bit_vector.select1(global_k).unwrap();
167                i = pos - offset;
168            } else {
169                let zeros_before = offset - self.ones_prefix[level];
170                let global_k = zeros_before + i;
171                let pos = self.bit_vector.select0(global_k).unwrap();
172                i = pos - offset;
173            }
174        }
175        Some(i)
176    }

Source

pub fn access(&self, k: usize) -> T

get k-th value

Source

pub fn rank(&self, val: T, range: Range<usize>) -> usize

the number of val in range

Examples found in repository ?

crates/library_checker/src/data_structure/static_range_frequency.rs (line 13)

6pub fn static_range_frequency(reader: impl Read, mut writer: impl Write) {
7    let s = read_all_unchecked(reader);
8    let mut scanner = Scanner::new(&s);
9    scan!(scanner, n, q, a: [usize; n]);
10    let wm = WaveletMatrix::new(a);
11    for _ in 0..q {
12        scan!(scanner, l, r, x: usize);
13        let ans = wm.rank(x, l..r);
14        iter_print!(writer, ans);
15    }
16}

Source

pub fn select(&self, val: T, k: usize) -> Option<usize>

index of k-th val

Source

pub fn quantile(&self, range: Range<usize>, k: usize) -> T

get k-th smallest value in range

Examples found in repository ?

crates/library_checker/src/data_structure/range_kth_smallest.rs (line 12)

6pub fn range_kth_smallest(reader: impl Read, mut writer: impl Write) {
7    let s = read_all_unchecked(reader);
8    let mut scanner = Scanner::new(&s);
9    scan!(scanner, n, q, a: [usize; n]);
10    let wm = WaveletMatrix::new(a);
11    for (l, r, k) in scanner.iter::<(usize, usize, usize)>().take(q) {
12        writeln!(writer, "{}", wm.quantile(l..r, k)).ok();
13    }
14}

Source

pub fn quantile_outer(&self, range: Range<usize>, k: usize) -> T

get k-th smallest value out of range

Source

pub fn rank_lessthan(&self, val: T, range: Range<usize>) -> usize

the number of value less than val in range

Examples found in repository ?

crates/competitive/src/data_structure/wavelet_matrix.rs (line 243)

242    pub fn rank_range(&self, valrange: Range<T>, range: Range<usize>) -> usize {
243        self.rank_lessthan(valrange.end, range.clone()) - self.rank_lessthan(valrange.start, range)
244    }

Source

pub fn rank_range(&self, valrange: Range<T>, range: Range<usize>) -> usize

the number of valrange in range

Source

pub fn query_less_than<F>(&self, val: T, range: Range<usize>, f: F)
where F: FnMut(usize, Range<usize>),

Source

pub fn build_fold<M>(&self, weights: &[M::T]) -> WaveletMatrixFold<'_, T, M>
where M: AbelianGroup,

Examples found in repository ?

crates/library_checker/src/data_structure/static_range_sum_with_upper_bound.rs (line 12)

6pub fn static_range_sum_with_upper_bound(reader: impl Read, mut writer: impl Write) {
7    let s = read_all_unchecked(reader);
8    let mut scanner = Scanner::new(&s);
9    scan!(scanner, n, q, a: [i64; n]);
10    let weights = a.clone();
11    let wm = WaveletMatrix::new(a);
12    let fold = wm.build_fold::<AdditiveOperation<i64>>(&weights);
13    for _ in 0..q {
14        scan!(scanner, l, r, x: i64);
15        let (count, sum) = fold.fold_lessthan_with_count(x + 1, l..r);
16        writeln!(writer, "{} {}", count, sum).ok();
17    }
18}