use std::cmp::max;
use std::marker::PhantomData;
use std::ops::Add;
pub trait PrefixOp<T> {
fn operation(t1: T, t2: T) -> T;
}
#[derive(Default, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug, Serialize, Deserialize)]
pub struct FenwickTree<T: Default + Ord, Op: PrefixOp<T>> {
tree: Vec<T>,
phantom: PhantomData<Op>,
}
impl<T: Ord + Default + Copy, Op: PrefixOp<T>> FenwickTree<T, Op> {
pub fn new(len: usize) -> FenwickTree<T, Op> {
FenwickTree {
tree: vec![T::default(); len + 1],
phantom: PhantomData,
}
}
pub fn get(&self, idx: usize) -> T {
let mut idx = idx + 1;
let mut sum = T::default();
while idx > 0 {
sum = Op::operation(sum, self.tree[idx]);
idx -= (idx as isize & -(idx as isize)) as usize;
}
sum
}
pub fn set(&mut self, idx: usize, val: T) {
let mut idx = idx + 1;
while idx < self.tree.len() {
self.tree[idx] = Op::operation(self.tree[idx], val);
idx += (idx as isize & -(idx as isize)) as usize;
}
}
}
#[derive(
Default, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug, Serialize, Deserialize,
)]
pub struct MaxOp;
impl<T: Copy + Ord> PrefixOp<T> for MaxOp {
fn operation(t1: T, t2: T) -> T {
max(t1, t2)
}
}
pub type MaxBitTree<T> = FenwickTree<T, MaxOp>;
#[derive(
Default, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug, Serialize, Deserialize,
)]
pub struct SumOp;
impl<T: Copy + Add> PrefixOp<T> for SumOp
where
T: Add<Output = T>,
{
fn operation(t1: T, t2: T) -> T {
t1 + t2
}
}
pub type SumBitTree<T> = FenwickTree<T, SumOp>;
#[cfg(test)]
mod test_bit_tree {
use super::MaxBitTree;
#[test]
pub fn test_bit_tree() {
let mut bit = MaxBitTree::new(10);
bit.set(0, (1, 0));
bit.set(1, (1, 1));
bit.set(2, (2, 2));
bit.set(3, (3, 3));
bit.set(4, (2, 4));
bit.set(5, (2, 5));
bit.set(6, (4, 6));
bit.set(7, (5, 7));
assert_eq!(bit.get(0), (1, 0));
assert_eq!(bit.get(1), (1, 1));
assert_eq!(bit.get(2), (2, 2));
assert_eq!(bit.get(3), (3, 3));
assert_eq!(bit.get(4), (3, 3));
assert_eq!(bit.get(5), (3, 3));
assert_eq!(bit.get(6), (4, 6));
assert_eq!(bit.get(7), (5, 7));
}
}