use proconio::{input, marker::Usize1}; #[derive(Clone)] pub struct AccumulateVector { accum_table: Vec, } impl AccumulateVector { pub fn new(v: &[usize]) -> Self { let mut accum_table = vec![0; v.len() + 1]; for i in 0..v.len() { accum_table[i + 1] = accum_table[i] + v[i]; } AccumulateVector { accum_table } } /// Returns: /// lenght of v pub fn len(&self) -> usize { self.accum_table.len() - 1 } /// Returns: /// sum(v[0..i]) pub fn rank(&self, i: usize) -> usize { self.accum_table[i] } } #[derive(Clone)] pub struct WaveletMatrix { bit_table: Vec, sorted_deduped_v: Vec, accum_table: Vec, accum_v: AccumulateVector, } impl WaveletMatrix { const WAVELET_MATRIX_HEIGHT: usize = 18; pub fn new(v: &[usize]) -> Self { let mut sorted_deduped_v = v.to_vec(); sorted_deduped_v.sort_unstable(); sorted_deduped_v.dedup(); let mut compress = v .iter() .map(|&x| sorted_deduped_v.partition_point(|&y| y < x)) .collect::>(); let mut bit_table = vec![]; let mut accum_table = vec![]; for i in (0..WaveletMatrix::WAVELET_MATRIX_HEIGHT).rev() { bit_table.push(AccumulateVector::new( &compress.iter().map(|&x| (x >> i) & 1).collect::>(), )); accum_table.push(AccumulateVector::new( &compress .iter() .map(|&x| { if (x >> i) & 1 == 0 { sorted_deduped_v[x] } else { 0 } }) .collect::>(), )); compress = compress .iter() .filter(|&x| (x >> i) & 1 == 0) .chain(compress.iter().filter(|&x| (x >> i) & 1 == 1)) .cloned() .collect::>(); } let accum_v = AccumulateVector::new(v); WaveletMatrix { bit_table, sorted_deduped_v, accum_table, accum_v, } } /// Returns: /// v[l..r].into_iter().filter(|y| y < upper).count() pub fn count_less_than(&self, mut l: usize, mut r: usize, upper: usize) -> usize { if r <= l { return 0; } let upper = self.sorted_deduped_v.partition_point(|&x| x < upper); let mut res = 0; for (i, bit) in (0..Self::WAVELET_MATRIX_HEIGHT) .rev() .zip(self.bit_table.iter()) { if (upper >> i) & 1 == 0 { l = l - bit.rank(l); r = r - bit.rank(r); } else { res += (r - l) - (bit.rank(r) - bit.rank(l)); l = bit.rank(l) + (bit.len() - bit.rank(bit.len())); r = bit.rank(r) + (bit.len() - bit.rank(bit.len())); } } res } /// Returns: /// v[l..r].into_iter().filter(|y| lower <= y ).count() pub fn count_more_than(&self, l: usize, r: usize, lower: usize) -> usize { if r <= l { return 0; } r - l - self.count_less_than(l, r, lower) } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().sorted().nth(k) /// } pub fn get_kth_smallest(&self, mut l: usize, mut r: usize, mut k: usize) -> Option { if r <= l || r - l <= k { return None; } let mut res = 0; for (i, bit) in (0..Self::WAVELET_MATRIX_HEIGHT) .rev() .zip(self.bit_table.iter()) { let j = (r - l) - (bit.rank(r) - bit.rank(l)); if k < j { l = l - bit.rank(l); r = r - bit.rank(r); } else { l = bit.rank(l) + (bit.len() - bit.rank(bit.len())); r = bit.rank(r) + (bit.len() - bit.rank(bit.len())); res |= 1 << i; k -= j; } } Some(self.sorted_deduped_v[res]) } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().sorted().rev().nth(k) /// } pub fn get_kth_largest(&self, l: usize, r: usize, k: usize) -> Option { if r <= l || r - l <= k { return None; } self.get_kth_smallest(l, r, r - l - 1 - k) } /// Returns: /// v[l..r].into_iter().filter(|y| lower <= y < upper).count() pub fn count(&self, l: usize, r: usize, lower: usize, upper: usize) -> usize { if r <= l { return 0; } self.count_less_than(l, r, upper) - self.count_less_than(l, r, lower) } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().filter(|x| x < upper).sum() /// } pub fn get_sum_less_than(&self, mut l: usize, mut r: usize, upper: usize) -> usize { if r <= l { return 0; } let upper = self.sorted_deduped_v.partition_point(|&x| x < upper); let mut res = 0; for (i, (bit, accum)) in (0..Self::WAVELET_MATRIX_HEIGHT) .rev() .zip(self.bit_table.iter().zip(self.accum_table.iter())) { if (upper >> i) & 1 == 0 { l = l - bit.rank(l); r = r - bit.rank(r); } else { res += accum.rank(r) - accum.rank(l); l = bit.rank(l) + (bit.len() - bit.rank(bit.len())); r = bit.rank(r) + (bit.len() - bit.rank(bit.len())); } } res } pub fn count_less_than_and_get_sum_less_than( &self, mut l: usize, mut r: usize, upper: usize, ) -> (usize, usize) { if r <= l { return (0, 0); } let upper = self.sorted_deduped_v.partition_point(|&x| x < upper); let mut count_less_than = 0; let mut sum_less_than = 0; for (i, (bit, accum)) in (0..Self::WAVELET_MATRIX_HEIGHT) .rev() .zip(self.bit_table.iter().zip(self.accum_table.iter())) { if (upper >> i) & 1 == 0 { l = l - bit.rank(l); r = r - bit.rank(r); } else { sum_less_than += accum.rank(r) - accum.rank(l); count_less_than += (r - l) - (bit.rank(r) - bit.rank(l)); l = bit.rank(l) + (bit.len() - bit.rank(bit.len())); r = bit.rank(r) + (bit.len() - bit.rank(bit.len())); } } (count_less_than, sum_less_than) } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().filter(|x| lower <= x).sum() /// } pub fn get_sum_more_than(&self, l: usize, r: usize, lower: usize) -> usize { if r <= l { return 0; } self.accum_v.rank(r) - self.accum_v.rank(l) - self.get_sum_less_than(l, r, lower) } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().sorted().take(k).sum() /// } pub fn get_sum_k_smallest(&self, mut l: usize, mut r: usize, mut k: usize) -> Option { if r <= l || r - l < k { return None; } let mut res = 0; let mut kth = 0; for (i, (bit, sum)) in (0..Self::WAVELET_MATRIX_HEIGHT) .rev() .zip(self.bit_table.iter().zip(self.accum_table.iter())) { let j = (r - l) - (bit.rank(r) - bit.rank(l)); if k < j { l = l - bit.rank(l); r = r - bit.rank(r); } else { res += sum.rank(r); res -= sum.rank(l); l = bit.rank(l) + (bit.len() - bit.rank(bit.len())); r = bit.rank(r) + (bit.len() - bit.rank(bit.len())); kth |= 1 << i; k -= j; } } res += k * kth; Some(res) } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().sorted().rev().take(k).sum() /// } pub fn get_sum_k_largest(&self, l: usize, r: usize, k: usize) -> Option { if r <= l || r - l < k { return None; } Some( self.accum_v.rank(r) - self.accum_v.rank(l) - self.get_sum_k_smallest(l, r, r - l - k).unwrap(), ) } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().filter(|x| lower <= x < upper).sum() /// } pub fn get_sum(&self, l: usize, r: usize, lower: usize, upper: usize) -> usize { if r <= l { return 0; } self.get_sum_less_than(l, r, upper) - self.get_sum_less_than(l, r, lower) } /// Returns: /// v[l..r].map(|y| y.abs_diff(x)).sum() pub fn get_sum_abs_diff(&self, l: usize, r: usize, x: usize) -> usize { let (c, sum_less_than) = self.count_less_than_and_get_sum_less_than(l, r, x); let d = r - l - c; let sum_more_than: usize = self.accum_v.rank(r) - self.accum_v.rank(l) - sum_less_than; (x * c - sum_less_than) + sum_more_than - x * d } } fn main() { input! { n: usize, q: usize, a: [isize; n], } let wm = WaveletMatrix::new( &a.iter() .cloned() .map(|x| (x + 1000000000) as usize) .collect::>(), ); let mut ans = vec![]; for _ in 0..q { input! { l: Usize1, r: usize, } let mid = wm.get_kth_largest(l, r, (r - l) / 2).unwrap(); ans.push(wm.get_sum_abs_diff(l, r, mid)); } println!( "{}", ans.into_iter() .map(|x| x.to_string()) .collect::>() .join("\n") ); }