Wavelet Matrix
struct WaveletMatrix {
n: usize,
lg: usize,
pref: Vec<Vec<usize>>,
}
impl WaveletMatrix {
fn from_vec(mut arr: Vec<usize>) -> Self {
let n = arr.len();
let mx = *arr.iter().max().unwrap();
let lg = mx.next_power_of_two().trailing_zeros() as usize + 1;
let mut bvs = vec![vec![]; lg];
for i in (0..lg).rev() {
bvs[i] = arr.iter().map(|x| (x >> i) & 1).collect();
let mut nxt = vec![];
nxt.extend((0..n).filter(|&j| bvs[i][j] == 0).map(|i| arr[i]));
nxt.extend((0..n).filter(|&j| bvs[i][j] == 1).map(|i| arr[i]));
arr = nxt;
}
Self {
n,
lg,
pref: (0..lg).map(|i| build(&bvs[i])).collect(),
}
}
// number of val in A[l..r]
fn count(&self, mut l: usize, mut r: usize, val: usize) -> usize {
for i in (0..self.lg).rev() {
if (val >> i) & 1 == 0 {
l = l - query(&self.pref[i], 0, l);
r = r - query(&self.pref[i], 0, r);
} else {
let p = self.n - query(&self.pref[i], 0, self.n);
l = p + query(&self.pref[i], 0, l);
r = p + query(&self.pref[i], 0, r);
}
}
r - l
}
// k-smallest (1-based) value in A[l..r]
fn k_smallest(&self, mut l: usize, mut r: usize, mut k: usize) -> usize {
let mut res = 0;
for i in (0..self.lg).rev() {
let cnt0 = (r - l) - query(&self.pref[i], l, r);
if k <= cnt0 {
l = l - query(&self.pref[i], 0, l);
r = r - query(&self.pref[i], 0, r);
} else {
let p = self.n - query(&self.pref[i], 0, self.n);
l = p + query(&self.pref[i], 0, l);
r = p + query(&self.pref[i], 0, r);
k -= cnt0;
res |= 1 << i;
}
}
res
}
}
fn build<T>(arr: &[T]) -> Vec<T>
where
T: Default + Copy + std::ops::Add<Output = T>,
{
let mut pref = vec![T::default(); arr.len()];
pref[0] = arr[0];
for i in 1..arr.len() {
pref[i] = pref[i - 1] + arr[i];
}
pref
}
// i..j
fn query<T>(pref: &[T], i: usize, j: usize) -> T
where
T: Default + Copy + std::ops::Sub<Output = T>,
{
if i == j {
T::default()
} else if i == 0 {
pref[j - 1]
} else {
pref[j - 1] - pref[i - 1]
}
}
Reference: