#[derive(Debug, Clone)]
pub struct CapacitySegmentTree {
n: usize,
tree: Vec<usize>,
}
impl CapacitySegmentTree {
pub fn new(max_capacity: usize) -> Self {
let mut n = 1;
while n < max_capacity + 1 {
n <<= 1;
}
let mut tree = vec![0usize; 2 * n];
tree[n - 1 + max_capacity] = max_capacity;
for i in (0..n - 1).rev() {
tree[i] = tree[2 * i + 1].max(tree[2 * i + 2]);
}
Self { n, tree }
}
pub fn new_empty(max_capacity: usize) -> Self {
let mut n = 1;
while n < max_capacity + 1 {
n <<= 1;
}
Self {
n,
tree: vec![0usize; 2 * n],
}
}
#[inline(always)]
pub fn update(&mut self, capacity: usize, value: usize) {
let mut idx = self.n - 1 + capacity;
self.tree[idx] = value;
while idx > 0 {
idx = (idx - 1) / 2;
let new_val = self.tree[2 * idx + 1].max(self.tree[2 * idx + 2]);
if self.tree[idx] == new_val {
break; }
self.tree[idx] = new_val;
}
}
#[inline(always)]
pub fn find_best_fit(&self, target: usize) -> Option<usize> {
if self.tree[0] < target {
return None;
}
let mut idx = 0;
while idx < self.n - 1 {
let left = 2 * idx + 1;
let right = 2 * idx + 2;
if self.tree[left] >= target {
idx = left; } else {
idx = right;
}
}
Some(idx - (self.n - 1))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_seeds_at_max_capacity() {
let tree = CapacitySegmentTree::new(10);
assert_eq!(tree.find_best_fit(1), Some(10));
assert_eq!(tree.find_best_fit(10), Some(10));
assert_eq!(tree.find_best_fit(11), None);
}
#[test]
fn test_new_empty_has_no_bins() {
let tree = CapacitySegmentTree::new_empty(10);
assert_eq!(tree.find_best_fit(1), None);
assert_eq!(tree.find_best_fit(10), None);
}
#[test]
fn test_update_and_query() {
let mut tree = CapacitySegmentTree::new_empty(10);
tree.update(7, 7);
assert_eq!(tree.find_best_fit(5), Some(7));
assert_eq!(tree.find_best_fit(7), Some(7));
assert_eq!(tree.find_best_fit(8), None);
tree.update(3, 3);
assert_eq!(tree.find_best_fit(2), Some(3)); assert_eq!(tree.find_best_fit(5), Some(7));
}
#[test]
fn test_best_fit_returns_smallest_qualifying() {
let mut tree = CapacitySegmentTree::new_empty(10);
tree.update(3, 3);
tree.update(5, 5);
tree.update(8, 8);
assert_eq!(tree.find_best_fit(4), Some(5));
assert_eq!(tree.find_best_fit(1), Some(3));
assert_eq!(tree.find_best_fit(6), Some(8));
}
#[test]
fn test_remove_capacity() {
let mut tree = CapacitySegmentTree::new(10);
tree.update(10, 0);
assert_eq!(tree.find_best_fit(1), None);
tree.update(5, 5);
assert_eq!(tree.find_best_fit(3), Some(5));
tree.update(5, 0);
assert_eq!(tree.find_best_fit(3), None);
}
#[test]
fn test_move_bin_between_capacities() {
let mut tree = CapacitySegmentTree::new(10);
tree.update(10, 0);
tree.update(7, 7);
assert_eq!(tree.find_best_fit(8), None);
assert_eq!(tree.find_best_fit(7), Some(7));
tree.update(7, 0);
tree.update(3, 3);
assert_eq!(tree.find_best_fit(4), None);
assert_eq!(tree.find_best_fit(3), Some(3));
}
#[test]
fn test_early_termination_correctness() {
let mut tree = CapacitySegmentTree::new_empty(8);
tree.update(5, 5);
tree.update(3, 3);
assert_eq!(tree.find_best_fit(4), Some(5));
tree.update(5, 0);
assert_eq!(tree.find_best_fit(4), None);
assert_eq!(tree.find_best_fit(3), Some(3));
}
#[test]
fn test_capacity_one() {
let mut tree = CapacitySegmentTree::new(1);
assert_eq!(tree.find_best_fit(1), Some(1));
assert_eq!(tree.find_best_fit(2), None);
tree.update(1, 0);
assert_eq!(tree.find_best_fit(1), None);
}
#[test]
fn test_power_of_two_capacity() {
let tree = CapacitySegmentTree::new(7);
assert_eq!(tree.find_best_fit(7), Some(7));
assert_eq!(tree.find_best_fit(8), None);
}
#[test]
fn test_multiple_updates_same_leaf() {
let mut tree = CapacitySegmentTree::new_empty(10);
tree.update(5, 5);
assert_eq!(tree.find_best_fit(5), Some(5));
tree.update(5, 0);
assert_eq!(tree.find_best_fit(5), None);
tree.update(5, 5);
assert_eq!(tree.find_best_fit(5), Some(5));
}
}