use proconio::{input, marker::Usize1}; #[derive(Clone)] pub struct WaveletMatrixRow { accum_table: Vec, } impl WaveletMatrixRow { pub fn new(bit: Vec) -> Self { let mut accum_table = vec![0; bit.len() + 1]; for i in 0..bit.len() { accum_table[i + 1] = accum_table[i] + bit[i]; } WaveletMatrixRow { accum_table } } pub fn len(&self) -> usize { self.accum_table.len() - 1 } pub fn access(&self, i: usize) -> usize { self.accum_table[i + 1] - self.accum_table[i] } pub fn rank(&self, i: usize) -> usize { self.accum_table[i] } } #[derive(Clone)] pub struct WaveletMatrix { accum: Vec, bit_table: Vec, accum_table: Vec, } impl WaveletMatrix { const WAVELET_MATRIX_HEIGHT: usize = 32; pub fn new(v: &[usize]) -> Self { let mut v = v.to_vec(); let mut accum = vec![0]; for i in 0..v.len() { accum.push(accum[i] + v[i]); } let mut bit_table = vec![]; let mut accum_table = vec![]; for i in (0..WaveletMatrix::WAVELET_MATRIX_HEIGHT).rev() { bit_table.push(WaveletMatrixRow::new( v.iter().map(|&x| (x >> i) & 1).collect(), )); accum_table.push(WaveletMatrixRow::new( v.iter() .map(|&x| if (x >> i) & 1 == 0 { x } else { 0 }) .collect::>(), )); // slow code // ``` // v.sort_by_key(|&x| (x >> i) & 1); // ``` v = v .iter() .filter(|&x| (x >> i) & 1 == 0) .chain(v.iter().filter(|&x| (x >> i) & 1 == 1)) .cloned() .collect::>(); } WaveletMatrix { accum, bit_table, accum_table, } } /// Returns: /// v[i] pub fn access(&self, mut i: usize) -> Option { if i >= self.bit_table[0].len() { return None; } let mut res = 0; for row in &self.bit_table { res <<= 1; res |= row.access(i); i = match row.access(i) { 0 => i - row.rank(i), 1 => row.len() - row.rank(row.len()) + row.rank(i), _ => { unreachable!() } }; } Some(res) } /// Returns: /// v[l..r].into_iter().filter(|y| **y == x).count() pub fn rank(&self, mut l: usize, mut r: usize, x: usize) -> usize { for (i, row) in self.bit_table.iter().rev().enumerate().rev() { if (x >> i) & 1 == 0 { l = l - row.rank(l); r = r - row.rank(r); } else { l = row.rank(l) + (row.len() - row.rank(row.len())); r = row.rank(r) + (row.len() - row.rank(row.len())); } } r - l } /// 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 l >= r || r - l <= k { return None; } let mut res = 0; for (i, row) in self.bit_table.iter().rev().enumerate().rev() { let j = (r - l) - (row.rank(r) - row.rank(l)); if j > k { l = l - row.rank(l); r = r - row.rank(r); } else { l = row.rank(l) + (row.len() - row.rank(row.len())); r = row.rank(r) + (row.len() - row.rank(row.len())); res |= 1 << i; k -= j; } } Some(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 + 1 + k) { None } else { self.get_kth_smallest(l, r, r - l - 1 - k) } } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().sorted().take(k).sum::() /// } pub fn get_sum_lower_k(&self, mut l: usize, mut r: usize, mut k: usize) -> usize { let mut res = 0; let mut kth = 0; for (i, (row, sum)) in (self.bit_table.iter().zip(self.accum_table.iter())) .rev() .enumerate() .rev() { let j = (r - l) - (row.rank(r) - row.rank(l)); if j > k { l = l - row.rank(l); r = r - row.rank(r); } else { res += sum.rank(r); res -= sum.rank(l); l = row.rank(l) + (row.len() - row.rank(row.len())); r = row.rank(r) + (row.len() - row.rank(row.len())); kth |= 1 << i; k -= j; } } res += k * kth; res } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().sorted().rev().take(k).sum::() /// } pub fn get_sum_upper_k(&self, l: usize, r: usize, k: usize) -> usize { self.accum[r] - self.accum[l] - self.get_sum_lower_k(l, r, r - l - k) } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().filter(|x| lower <= *x && *x < upper).sum::() /// } pub fn get_sum(&self, l: usize, r: usize, lower: usize, upper: usize) -> usize { self.get_sum_lower_k(l, r, self.count(l, r, 0, upper)) - self.get_sum_lower_k(l, r, self.count(l, r, 0, lower)) } /// 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 { let mut res = 0; for (i, row) in self.bit_table.iter().rev().enumerate().rev() { if (upper >> i) & 1 == 0 { l = l - row.rank(l); r = r - row.rank(r); } else { res += (r - l) - (row.rank(r) - row.rank(l)); l = row.rank(l) + (row.len() - row.rank(row.len())); r = row.rank(r) + (row.len() - row.rank(row.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 { r - l - self.count_less_than(l, r, lower) } /// Returns: /// v[l..r].into_iter().filter(|y| lower <= **y && **y < upper).count() pub fn count(&self, l: usize, r: usize, lower: usize, upper: usize) -> usize { self.count_less_than(l, r, upper) - self.count_less_than(l, r, lower) } /// Returns: /// { /// use itertools::Itertools; /// v[l..r].into_iter().filter(|y| lower <= **y).sorted().nth(k) /// } pub fn get_above_kth_smallest( &self, l: usize, r: usize, lower: usize, k: usize, ) -> Option { let cnt = self.count(l, r, 0, lower); if cnt + k >= r - l { None } else { Some(self.get_kth_smallest(l, r, cnt + k).unwrap()) } } /// Returns: /// v[l..r].into_iter().filter(|y| **y < upper).sorted().rev().nth(k) pub fn get_below_kth_largest( &self, l: usize, r: usize, upper: usize, k: usize, ) -> Option { let cnt = self.count(l, r, 0, upper); if cnt <= k { None } else { Some(self.get_kth_smallest(l, r, cnt - 1 - k).unwrap()) } } /// 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 = self.count_less_than(l, r, x); let d = self.count_more_than(l, r, x); x * c - self.get_sum_lower_k(l, r, c) + self.get_sum_upper_k(l, r, d) - x * d } } fn main() { input! { n: usize, q: usize, a: [isize; n], } let wm = WaveletMatrix::new( &a.iter() .cloned() .map(|x| 2 * (x + 1000000000) as usize) .collect::>(), ); let mut ans = vec![]; for _ in 0..q { input! { l: Usize1, r: usize, } if (r - l) % 2 == 1 { let mid = wm.get_kth_largest(l, r, (r - l) / 2).unwrap(); ans.push(wm.get_sum_abs_diff(l, r, mid) / 2); } else { let mid0 = wm.get_kth_largest(l, r, (r - l) / 2 - 1).unwrap(); let mid1 = wm.get_kth_largest(l, r, (r - l) / 2).unwrap(); ans.push(wm.get_sum_abs_diff(l, r, (mid0 + mid1) / 2) / 2); } } println!( "{}", ans.into_iter() .map(|x| x.to_string()) .collect::>() .join("\n") ); }