#![no_std]
#![warn(missing_docs)]
extern crate alloc;
use core::cmp::{min, max};
use alloc::vec::Vec;
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct ChthollyNode {
left: usize,
right: usize,
value: u64,
}
impl ChthollyNode {
pub fn contains(&self, x: usize) -> bool {
self.left <= x && x <= self.right
}
pub fn len(&self) -> usize {
self.right - self.left + 1
}
}
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct ChthollyTree(Vec<ChthollyNode>);
impl ChthollyTree {
pub fn from_slice(data: &[u64]) -> Self {
Self(data.iter().enumerate().map(|(i, d)| {
ChthollyNode {
left: i,
right: i,
value: *d,
}
}).collect())
}
pub fn split(&mut self, middle: usize) -> Option<&ChthollyNode> {
match self.split_inner(middle) {
Some(index) => Some(&self.0[index]),
None => None,
}
}
pub fn merge(&mut self, left: usize, right: usize, value: u64) {
let index = self.split_inner(left);
let end = self.split_inner(right + 1);
match index {
Some(index) => {
self.0[index].value = value;
self.0[index].right = right;
debug_assert!(self.0[index].left <= self.0[index].right);
match end {
None => {
while index + 1 < self.0.len() && self.0[index + 1].left <= right {
self.0.remove(index + 1);
}
},
Some(end) => {
self.0.drain((index + 1)..end);
},
}
debug_assert!(index + 1 >= self.0.len() ||
self.0[index + 1].left == self.0[index].right + 1);
},
None => {
self.0.push(ChthollyNode { left, right, value });
self.sort_inner();
},
}
}
pub fn add(&mut self, left: usize, right: usize, value: u64) {
self.split_inner(right + 1);
let start = match self.split_inner(left) {
Some(start) => start,
None => return,
};
for index in start..self.0.len() {
if self.0[index].left > right {
break
}
self.0[index].value += value;
}
}
pub fn nth(&self, left: usize, right: usize, mut n: usize) -> Option<u64> {
let mut index = match self.0.binary_search_by(|node| {
node.left.cmp(&left)
}) {
Ok(index) => index,
Err(index) => {
if index > 0 {
index - 1
} else {
return None
}
},
};
let mut values = Vec::new();
loop {
if index >= self.0.len() || self.0[index].left > right {
break
}
let left = max(left, self.0[index].left);
let right = min(right, self.0[index].right);
let len = right - left + 1;
values.push((len, self.0[index].value));
index += 1;
}
values.sort_unstable_by_key(|(_, v)| *v);
let mut target = 0;
loop {
if target >= values.len() {
return None
}
if n < values[target].0 || n == 0 {
return Some(values[target].1)
}
n -= values[target].0;
target += 1;
}
}
pub fn pow_sum(&self, left: usize, right: usize, power: u32, modulo: u64) -> u64 {
let mut index = match self.0.binary_search_by(|node| {
node.left.cmp(&left)
}) {
Ok(index) => index,
Err(index) => {
if index > 0 {
index - 1
} else {
return 0
}
},
};
let mut sum = 0;
loop {
if index >= self.0.len() || self.0[index].left > right {
break
}
let left = max(left, self.0[index].left);
let right = min(right, self.0[index].right);
let len = right - left + 1;
sum = (sum + len as u64 * mod_pow(self.0[index].value, power, modulo)) % modulo;
index += 1;
}
sum
}
fn sort_inner(&mut self) {
self.0.sort_unstable_by_key(|node| node.left);
}
fn split_inner(&mut self, middle: usize) -> Option<usize> {
let index = match self.0.binary_search_by(|node| {
node.left.cmp(&middle)
}) {
Ok(index) => index,
Err(index) => {
if index > 0 {
index - 1
} else {
return None
}
},
};
if self.0[index].left == middle {
return Some(index)
}
if self.0[index].right < middle {
return None
}
let new = ChthollyNode {
left: middle,
right: self.0[index].right,
value: self.0[index].value,
};
debug_assert!(new.left <= new.right);
self.0.insert(index + 1, new);
self.0[index].right = middle - 1;
debug_assert!(index + 1 < self.0.len() ||
self.0[index + 1].left == self.0[index].right + 1);
debug_assert!(self.0[index].left <= self.0[index].right);
Some(index + 1)
}
}
fn mod_pow(mut x: u64, mut n: u32, modulo: u64) -> u64 {
x = x % modulo;
let mut ret = 1 % modulo;
while n > 0 {
if n & 1 == 1 {
ret = ret * x % modulo;
}
x = x * x % modulo;
n = n >> 1;
}
ret
}
#[cfg(test)]
#[allow(unused)]
mod tests {
use alloc::vec;
use super::*;
const V_MAX_BOUND: usize = 1_000_000_000;
const SEED_MAX: usize = 1_000_000_007;
struct CF896CRng(usize);
impl CF896CRng {
fn next(&mut self) -> usize {
let ret = self.0;
self.0 = (self.0 * 7 + 13) % SEED_MAX;
ret
}
}
#[derive(Clone)]
enum Op {
Add(usize, usize, u64),
Assign(usize, usize, u64),
Nth(usize, usize, usize),
PowSum(usize, usize, u32, u64),
}
fn random_array(n: usize, vmax: usize, rng: &mut CF896CRng) -> Vec<u64> {
debug_assert!(vmax <= V_MAX_BOUND);
let mut ret = Vec::new();
for _ in 0..n {
ret.push((rng.next() % vmax + 1) as u64);
}
ret
}
fn random_ops(n: usize, m: usize, vmax: usize, rng: &mut CF896CRng) -> Vec<Op> {
debug_assert!(vmax <= V_MAX_BOUND);
let mut ret = Vec::new();
for _ in 0..m {
let opi = rng.next() % 4 + 1;
let (l, r) = {
let l = rng.next() % n + 1;
let r = rng.next() % n + 1;
if l > r {
(r, l)
} else {
(l, r)
}
};
let x = if opi == 3 {
rng.next() % (r - l + 1) + 1
} else {
rng.next() % vmax + 1
};
let op = match opi {
1 => Op::Add(l, r, x as u64),
2 => Op::Assign(l, r, x as u64),
3 => Op::Nth(l, r, x),
4 => {
let y = rng.next() % vmax + 1;
Op::PowSum(l, r, x as u32, y as u64)
},
_ => unreachable!("opi is modulo 4 plus 1"),
};
ret.push(op);
}
ret
}
fn test_vector(n: usize, m: usize, seed: usize, vmax: usize, expected: Vec<u64>) {
let mut rng = CF896CRng(seed);
let array = random_array(n, vmax, &mut rng);
let ops = random_ops(n, m, vmax, &mut rng);
let mut tree = ChthollyTree::from_slice(&array);
let mut output = Vec::new();
for op in ops.clone() {
match op {
Op::Add(l, r, x) => {
tree.add(l - 1 , r - 1, x);
},
Op::Assign(l, r, x) => {
tree.merge(l - 1, r - 1, x);
},
Op::Nth(l, r, x) => {
let n = tree.nth(l - 1, r - 1, x - 1).expect("Vector test failed to find n");
output.push(n);
},
Op::PowSum(l, r, x, y) => {
let z = tree.pow_sum(l - 1, r - 1, x as u32, y as u64);
output.push(z);
},
}
}
assert_eq!(output, expected);
}
#[test]
fn vector1() {
test_vector(10, 10, 7, 9, vec![2, 1, 0, 3]);
}
#[test]
fn vector2() {
test_vector(10, 10, 9, 9, vec![1, 1, 3, 3]);
}
}